server/modules/hf_dataset_loader.py

import os
import time
import uuid
from typing import Optional

from dotenv import load_dotenv
from tqdm.auto import tqdm
from datasets import load_dataset
from langchain_huggingface import HuggingFaceEmbeddings
from modules.load_vectorstore import get_index
from core.settings import get_settings

load_dotenv()

_BATCH_SIZE = 100


# ---------------------------------------------------------------------------
# Normalizers — each converts a raw HF row into {text, metadata}
# ---------------------------------------------------------------------------

def _norm_pubmedqa(row: dict) -> Optional[dict]:
    ctx = row.get("context", {})
    passages = ctx.get("contexts", []) if isinstance(ctx, dict) else []
    text = " ".join(
        filter(None, [row.get("question", ""), " ".join(passages), row.get("long_answer", "")])
    ).strip()
    if not text:
        return None
    return {
        "text": text,
        "metadata": {
            "source": "pubmedqa",
            "pubid": str(row.get("pubid", "")),
            "label": str(row.get("final_decision", "")),
        },
    }


def _norm_mental_health(row: dict) -> Optional[dict]:
    text = f"Patient: {row.get('Context', '')}\nCounselor: {row.get('Response', '')}".strip()
    return {"text": text, "metadata": {"source": "mental_health_counseling"}} if text else None


def _norm_mediqa(row: dict) -> Optional[dict]:
    text = " ".join(
        filter(None, [row.get("instruction"), row.get("input"), row.get("output")])
    ).strip()
    return {"text": text, "metadata": {"source": "medical_meadow_mediqa"}} if text else None


def _norm_medqa_usmle(row: dict) -> Optional[dict]:
    options = " | ".join(
        filter(None, [row.get(f"ending{i}") for i in range(4)])
    )
    text = f"{row.get('sent1', '')} Options: {options}".strip()
    return {
        "text": text,
        "metadata": {"source": "medqa_usmle", "answer_idx": str(row.get("label", ""))},
    } if text else None


# ---------------------------------------------------------------------------
# Registry
# ---------------------------------------------------------------------------

DATASET_REGISTRY: dict[str, dict] = {
    "pubmedqa": {
        "hf_id": "qiaojin/PubMedQA",
        "config": "pqa_labeled",
        "split": "train",
        "normalizer": _norm_pubmedqa,
        "description": "PubMedQA — biomedical Q&A backed by PubMed abstracts (MIT license)",
    },
    "mental_health": {
        "hf_id": "Amod/mental_health_counseling_conversations",
        "config": None,
        "split": "train",
        "normalizer": _norm_mental_health,
        "description": "Mental health counseling conversations with professional responses",
    },
    "mediqa": {
        "hf_id": "medalpaca/medical_meadow_mediqa",
        "config": None,
        "split": "train",
        "normalizer": _norm_mediqa,
        "description": "Medical Meadow MediQA — curated clinical Q&A (Nature Scientific Data)",
    },
    "medqa_usmle": {
        "hf_id": "GBaker/MedQA-USMLE-4-options-hf",
        "config": None,
        "split": "train",
        "normalizer": _norm_medqa_usmle,
        "description": "MedQA-USMLE — US medical licensing exam questions (CC-BY-SA-4.0)",
    },
}


def list_available_datasets() -> list[dict]:
    return [{"name": k, "description": v["description"]} for k, v in DATASET_REGISTRY.items()]


# ---------------------------------------------------------------------------
# Ingestion
# ---------------------------------------------------------------------------

def load_hf_datasets_to_pinecone(
    dataset_names: list[str],
    max_samples: int = 500,
) -> dict[str, dict]:
    settings = get_settings()
    if settings.huggingface_hub_token:
        os.environ["HUGGINGFACE_HUB_TOKEN"] = settings.huggingface_hub_token

    embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
    index = get_index()
    results: dict[str, dict] = {}

    for name in dataset_names:
        if name not in DATASET_REGISTRY:
            results[name] = {"status": "error", "detail": f"Unknown dataset '{name}'", "records_upserted": 0}
            continue

        cfg = DATASET_REGISTRY[name]

        try:
            ds = load_dataset(cfg["hf_id"], cfg["config"], split=cfg["split"], trust_remote_code=True)
        except Exception as exc:
            results[name] = {"status": "error", "detail": str(exc), "records_upserted": 0}
            continue

        if max_samples < len(ds):
            ds = ds.select(range(max_samples))

        records = [r for row in ds if (r := cfg["normalizer"](row))]

        if not records:
            results[name] = {"status": "skipped", "detail": "No valid records after normalisation", "records_upserted": 0}
            continue

        upserted = 0
        for i in tqdm(range(0, len(records), _BATCH_SIZE), desc=f"Upserting {name}"):
            batch = records[i : i + _BATCH_SIZE]
            texts = [r["text"] for r in batch]
            metas = [r["metadata"] for r in batch]

            embeddings = embed_model.embed_documents(texts)
            vectors = [
                {
                    "id": f"{name}-{i + j}-{uuid.uuid4().hex[:8]}",
                    "values": embeddings[j],
                    "metadata": {**metas[j], "text": texts[j]},
                }
                for j in range(len(batch))
            ]
            index.upsert(vectors=vectors)
            upserted += len(vectors)
            time.sleep(0.1)

        results[name] = {"status": "success", "records_upserted": upserted, "detail": ""}

    return results