config.yaml

# SafeRAG Configuration File

# Model Configuration
models:
  embedding:
    name: "BAAI/bge-large-en-v1.5"
    device: "cuda"
    batch_size: 32
  
  reranker:
    name: "cross-encoder/ms-marco-MiniLM-L-6-v2"
    device: "cuda"
    batch_size: 32
  
  generator:
    name: "openai/gpt-oss-20b"
    tensor_parallel_size: 1
    gpu_memory_utilization: 0.9
    max_tokens: 512
    temperature: 0.7
    top_p: 0.9
  
  calibration:
    type: "logistic"  # logistic, random_forest, mlp
    input_dim: 16
    hidden_dim: 64

# Data Configuration
data:
  datasets:
    - "hotpotqa"
    - "triviaqa"
    - "nq_open"
  
  knowledge_base:
    name: "wikipedia"
    language: "en"
    date: "20231101"
  
  preprocessing:
    max_sentence_length: 512
    min_sentence_length: 20
    cache_dir: "./cache"

# Index Configuration
index:
  type: "ivf"  # flat, ivf
  dimension: 1024
  nlist: 4096
  save_path: "./index/safrag"

# Retrieval Configuration
retrieval:
  k: 20
  rerank_k: 10
  batch_size: 32
  similarity_threshold: 0.3

# Risk Calibration Configuration
calibration:
  tau1: 0.3  # Low risk threshold
  tau2: 0.7  # High risk threshold
  
  features:
    - "num_passages"
    - "avg_similarity"
    - "std_similarity"
    - "max_similarity"
    - "min_similarity"
    - "score_variance"
    - "avg_token_overlap"
    - "max_token_overlap"
    - "avg_entity_overlap"
    - "max_entity_overlap"
    - "passage_consistency"
    - "passage_consistency_std"
    - "min_passage_similarity"
    - "diversity"
    - "topic_variance"

# Evaluation Configuration
evaluation:
  metrics:
    qa:
      - "exact_match"
      - "f1"
      - "rouge1"
      - "rouge2"
      - "rougeL"
    
    attribution:
      - "precision"
      - "recall"
      - "f1"
      - "citation_coverage"
      - "citation_accuracy"
    
    calibration:
      - "ece"
      - "mce"
      - "auroc"
      - "auprc"
    
    system:
      - "throughput"
      - "latency"
      - "gpu_utilization"
      - "memory_usage"
  
  test_size: 0.2
  random_state: 42
  cv_folds: 5

# System Configuration
system:
  device: "cuda"
  num_workers: 4
  batch_size: 32
  max_memory_gb: 16
  
  monitoring:
    enabled: true
    interval: 1  # seconds
    metrics:
      - "cpu"
      - "memory"
      - "gpu"
      - "disk"

# Output Configuration
output:
  results_dir: "./results"
  logs_dir: "./logs"
  models_dir: "./models"
  plots_dir: "./plots"
  
  formats:
    - "json"
    - "csv"
    - "html"
  
  save_predictions: true
  save_features: true
  save_plots: true

# Logging Configuration
logging:
  level: "INFO"
  format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
  file: "./logs/safrag.log"
  max_size: "10MB"
  backup_count: 5

# Hugging Face Configuration
huggingface:
  cache_dir: "./cache"
  token: null  # Set your HF token here
  hub_url: "https://huggingface.co"
  
  spaces:
    app_name: "safrag-demo"
    hardware: "cpu"  # cpu, gpu, cpu-basic, gpu-basic
    visibility: "public"

# Experiment Configuration
experiments:
  baseline:
    enabled: true
    output_dir: "./results/baseline"
  
  safrag:
    enabled: true
    output_dir: "./results/safrag"
  
  ablation:
    enabled: true
    output_dir: "./results/ablation"
    
    studies:
      - "no_reranking"
      - "no_calibration"
      - "different_embeddings"
      - "different_thresholds"
      - "different_calibration_models"
      - "different_retrieval_k"
  
  comprehensive:
    enabled: true
    output_dir: "./results/comprehensive"