export_dataset.py

"""
export_dataset.py
=================
Esporta tutti i dati da ChromaDB + seed_italian.py in formato JSONL
compatibile con HuggingFace (ChatML / SFTTrainer).

Output: dataset.jsonl  (pronto per AutoTrain o finetune.py)

Uso:
    python export_dataset.py
    python export_dataset.py --out mio_dataset.jsonl --min-chars 50
"""

import argparse
import json
import os
import sys
from pathlib import Path

SYSTEM_PROMPT = (
    "Sei GenerAI, un assistente AI specializzato in lingua italiana. "
    "Rispondi in modo chiaro, preciso e sempre in italiano."
)


def _row(question: str, answer: str) -> dict:
    """Formato ChatML — compatibile con SFTTrainer e HF AutoTrain."""
    return {
        "messages": [
            {"role": "system",    "content": SYSTEM_PROMPT},
            {"role": "user",      "content": question.strip()},
            {"role": "assistant", "content": answer.strip()},
        ]
    }


def load_from_seed() -> list[dict]:
    """Carica le regole di grammatica da seed_italian.py."""
    from seed_italian import GRAMMAR_SEED
    rows = [_row(item["q"], item["a"]) for item in GRAMMAR_SEED]
    print(f"[seed]    {len(rows)} esempi caricati da seed_italian.py")
    return rows


def load_from_chromadb(min_chars: int = 80) -> list[dict]:
    """Carica i documenti archiviati da ChromaDB (domande web salvate)."""
    try:
        import chromadb
        from chromadb.utils.embedding_functions import SentenceTransformerEmbeddingFunction
    except ImportError:
        print("[chromadb] chromadb non installato — saltato.")
        return []

    db_path = "./database"
    if not os.path.exists(db_path):
        print("[chromadb] Nessun database trovato — saltato.")
        return []

    ef = SentenceTransformerEmbeddingFunction(
        model_name="paraphrase-multilingual-MiniLM-L12-v2"
    )
    client = chromadb.PersistentClient(path=db_path)
    try:
        col = client.get_collection(name="generai", embedding_function=ef)
    except Exception:
        print("[chromadb] Collezione 'generai' non trovata — saltato.")
        return []

    total = col.count()
    if total == 0:
        print("[chromadb] Database vuoto — saltato.")
        return []

    results = col.get(include=["metadatas"], limit=total)
    rows = []
    skipped = 0
    for meta in results["metadatas"]:
        source = meta.get("source", "")
        if source == "grammatica_italiana":
            continue  # già in seed, evita duplicati
        q = meta.get("query", "").strip()
        a = meta.get("answer", "").strip()
        if not q or not a or len(a) < min_chars:
            skipped += 1
            continue
        rows.append(_row(q, a))

    print(f"[chromadb] {len(rows)} esempi caricati ({skipped} saltati per qualità).")
    return rows


def deduplicate(rows: list[dict]) -> list[dict]:
    seen = set()
    out = []
    for row in rows:
        key = row["messages"][1]["content"][:80].lower()
        if key not in seen:
            seen.add(key)
            out.append(row)
    return out


def main():
    parser = argparse.ArgumentParser(description="Esporta dataset per HuggingFace fine-tuning")
    parser.add_argument("--out",       default="dataset.jsonl", help="File di output (default: dataset.jsonl)")
    parser.add_argument("--min-chars", type=int, default=80,    help="Lunghezza minima risposta (default: 80)")
    parser.add_argument("--no-seed",   action="store_true",     help="Non includere seed_italian.py")
    parser.add_argument("--no-db",     action="store_true",     help="Non includere ChromaDB")
    args = parser.parse_args()

    rows = []
    if not args.no_seed:
        rows += load_from_seed()
    if not args.no_db:
        rows += load_from_chromadb(min_chars=args.min_chars)

    rows = deduplicate(rows)

    if not rows:
        print("❌ Nessun dato trovato. Aggiungi dati alla KB prima di esportare.")
        sys.exit(1)

    out_path = Path(args.out)
    with open(out_path, "w", encoding="utf-8") as f:
        for row in rows:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")

    print(f"\n✅ Dataset esportato: {out_path.resolve()}")
    print(f"   Totale esempi : {len(rows)}")
    print(f"   Formato       : ChatML (messages: system/user/assistant)")
    print(f"\nProssimo passo:")
    print(f"  → Fine-tuning locale : python finetune.py --dataset {args.out}")
    print(f"  → HuggingFace AutoTrain: carica {args.out} su https://huggingface.co/autotrain")


if __name__ == "__main__":
    main()