candding 0.1.0

use candle_core::DType;
use candle_core::Device;

use super::*;
use crate::embed::builder::reject_unsupported_dtype;
use crate::error::CanddingError;
use crate::models::RerankHead;
use crate::models::Unembedding;
use crate::models::jina_bert::AlibiSlopeVariant;
use crate::models::xlm_roberta_flash::PositionScheme;
use crate::registry;
use crate::registry::Backend;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PairTruncationPolicy;

fn row(cpu: Support, metal: Support, cuda: Support) -> BackendSupport {
  BackendSupport { cpu, metal, cuda }
}

fn axis() -> DtypeBackendSupport {
  DtypeBackendSupport {
    f32: row(Support::Verified, Support::Verified, Support::Untested),
    f16: row(Support::Unsupported, Support::Verified, Support::Untested),
    bf16: row(
      Support::Unsupported,
      Support::Unsupported,
      Support::Unsupported,
    ),
  }
}

#[test]
fn a_cpu_only_unsupported_cell_is_refused_on_the_cpu_as_accelerator_only() {
  let err = reject_unsupported_dtype("m", DType::F16, axis(), &Device::Cpu).unwrap_err();
  assert!(
    matches!(
      err,
      CanddingError::UnsupportedOnDevice {
        backend: Backend::Cpu,
        dtype: DType::F16,
        ..
      }
    ),
    "{err}"
  );
  let message = err.to_string();
  assert!(message.contains("accelerator-only"), "{message}");
  assert!(!message.contains("NaN"), "{message}");
}

#[test]
fn a_dtype_unsupported_on_every_device_is_refused_as_unsupported() {
  let err = reject_unsupported_dtype("m", DType::BF16, axis(), &Device::Cpu).unwrap_err();
  assert!(
    matches!(&err, CanddingError::UnsupportedDtype { model, dtype: DType::BF16 } if model == "m"),
    "{err}"
  );
}

#[test]
fn a_verified_cell_passes() {
  assert!(reject_unsupported_dtype("m", DType::F32, axis(), &Device::Cpu).is_ok());
}

#[test]
fn bge_reranker_base_is_refused_through_its_registered_rows() {
  let descriptor = registry::rerank::lookup("BAAI/bge-reranker-base").expect("registered");
  let rows = descriptor.dtype_backends;
  let f16 = reject_unsupported_dtype("BAAI/bge-reranker-base", DType::F16, rows, &Device::Cpu);
  assert!(matches!(
    f16,
    Err(CanddingError::UnsupportedOnDevice {
      backend: Backend::Cpu,
      dtype: DType::F16,
      ..
    })
  ));
  let bf16 = reject_unsupported_dtype("BAAI/bge-reranker-base", DType::BF16, rows, &Device::Cpu);
  assert!(matches!(
    bf16,
    Err(CanddingError::UnsupportedOnDevice {
      backend: Backend::Cpu,
      dtype: DType::BF16,
      ..
    })
  ));
}

/// Every convention differs from a hardcoded default, so a constant in place of any field fails its assertion.
#[test]
fn conventions_from_descriptor_carries_every_field_without_hardcoding_one() {
  let descriptor = RerankDescriptor {
    id: "org/stub".to_string(),
    family: Family::JinaBert,
    head: RerankHead::RobertaClassificationHead,
    default_max_length: 512,
    score_activation: ScoreActivation::Sigmoid,
    input: RerankInput::Pair {
      truncation: PairTruncationPolicy::PreserveQueryToThreeQuarters,
    },
    alibi_variant: Some(AlibiSlopeVariant::HalvedInterpolatedTail),
    position_scheme: Some(PositionScheme::Absolute),
    default_dtype: DType::F32,
    backends: BackendSupport::untested(),
    dtype_backends: DtypeBackendSupport::untested(),
    license: "unknown".to_string(),
    size_gb: 0.0,
    phase: u8::MAX,
    weights: WeightsFormat::Safetensors,
    gated: false,
  };
  let conventions = conventions_from_descriptor(&descriptor);
  assert_eq!(
    conventions.head,
    RerankHead::RobertaClassificationHead,
    "head must carry through"
  );
  assert_eq!(
    conventions.score_activation,
    ScoreActivation::Sigmoid,
    "score_activation must carry through"
  );
  assert_eq!(
    conventions.alibi_variant,
    Some(AlibiSlopeVariant::HalvedInterpolatedTail),
    "alibi_variant must carry through"
  );
  assert_eq!(
    conventions.position_scheme,
    Some(PositionScheme::Absolute),
    "position_scheme must carry through"
  );
}

/// No flash checkpoint's `config.json` gives a `model_type`, so the flash key dispatches without reading it.
#[test]
fn a_flash_entry_dispatches_by_its_family_without_reading_config_json() {
  let mut descriptor = registry::rerank::lookup("BAAI/bge-reranker-base").expect("registered");
  let missing = std::path::Path::new("/nonexistent/candding/config.json");
  descriptor.family = Family::XlmRobertaFlash;
  assert_eq!(
    architecture(&descriptor, missing).unwrap(),
    xlm_roberta_flash::ARCHITECTURE
  );
  assert_eq!(
    serde_json::to_value(Family::XlmRobertaFlash).unwrap(),
    xlm_roberta_flash::ARCHITECTURE
  );
  descriptor.family = Family::XlmRoberta;
  assert!(matches!(
    architecture(&descriptor, missing),
    Err(CanddingError::Io(_))
  ));
  let dir = tempfile::tempdir().unwrap();
  let config = dir.path().join("config.json");
  std::fs::write(&config, r#"{"model_type": "xlm-roberta"}"#).unwrap();
  assert_eq!(architecture(&descriptor, &config).unwrap(), "xlm-roberta");
}

#[test]
fn the_registered_jina_reranker_maps_onto_its_own_conventions() {
  let descriptor = registry::rerank::lookup("jinaai/jina-reranker-v1-tiny-en").expect("registered");
  assert_eq!(
    conventions_from_descriptor(&descriptor),
    RerankConventions {
      head: RerankHead::BertPooler,
      score_activation: ScoreActivation::Sigmoid,
      alibi_variant: Some(AlibiSlopeVariant::HalvedInterpolatedTail),
      position_scheme: None,
    }
  );
}

#[test]
fn the_registered_flash_reranker_maps_onto_its_own_conventions() {
  let descriptor =
    registry::rerank::lookup("jinaai/jina-reranker-v2-base-multilingual").expect("registered");
  assert_eq!(
    conventions_from_descriptor(&descriptor),
    RerankConventions {
      head: RerankHead::RobertaClassificationHead,
      score_activation: ScoreActivation::Sigmoid,
      alibi_variant: None,
      position_scheme: Some(PositionScheme::Absolute),
    }
  );
}

#[test]
fn conventions_from_descriptor_carries_the_logit_score_ids_and_tie() {
  let registered =
    registry::rerank::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").expect("registered");
  for (true_token_id, false_token_id) in [(9693, 2152), (2152, 9693), (7, 11)] {
    for unembedding in [Unembedding::Tied, Unembedding::Untied] {
      let head = RerankHead::LogitScore {
        true_token_id,
        false_token_id,
        unembedding,
      };
      let descriptor = RerankDescriptor {
        head,
        ..registered.clone()
      };
      assert_eq!(conventions_from_descriptor(&descriptor).head, head);
    }
  }
}