Rag carti.py

Rag carti.py

"""
Sistem RAG peste cărțile tale: transformă paragrafele în vectori, caută semantic
la o întrebare modernă și generează un răspuns bazat DOAR pe textele găsite.

Instalare:
    pip install "sentence-transformers>=3.0" "chromadb>=0.5" \
                "transformers>=4.44" accelerate torch

    Model de embedding recomandat pentru română: intfloat/multilingual-e5-large
    (dacă vrei ceva mai ușor: paraphrase-multilingual-MiniLM-L12-v2)

Rulare:
    python rag_carti.py --index          # o singură dată: construiește baza de vectori
    python rag_carti.py --ask "..."      # pune întrebări
"""

import os
import re
import glob
import argparse
import torch
import chromadb
from sentence_transformers import SentenceTransformer
from transformers import AutoModelForCausalLM, AutoTokenizer

# ----------------------------------------------------------------------
# CONFIGURARE
# ----------------------------------------------------------------------
TEXTS_DIR = "./carti_txt"                       # fișierele .txt ale cărților
DB_DIR = "./chroma_carti"                        # unde se salvează baza de vectori
COLLECTION = "carti_vechi"
EMB_MODEL = "intfloat/multilingual-e5-large"
GEN_MODEL = "mistralai/Mistral-7B-Instruct-v0.3" # modelul care formulează răspunsul
CHUNK_WORDS = 220                                # ~un paragraf per fragment
CHUNK_OVERLAP = 40                               # cuvinte de suprapunere între fragmente
TOP_K = 5                                        # câte fragmente se aduc la fiecare întrebare


# ----------------------------------------------------------------------
# 1. SEGMENTARE ÎN FRAGMENTE (chunking) cu suprapunere
# ----------------------------------------------------------------------
def fragmenteaza(text, marime, suprapunere):
    cuvinte = text.split()
    pas = marime - suprapunere
    for start in range(0, len(cuvinte), pas):
        bucata = cuvinte[start:start + marime]
        if len(bucata) < 30:
            continue
        yield " ".join(bucata)


def curata(text):
    text = re.sub(r"-\n", "", text)             # unește cuvintele despărțite la capăt de rând
    text = re.sub(r"\s+", " ", text)
    return text.strip()


# ----------------------------------------------------------------------
# 2. INDEXARE: text -> vectori -> ChromaDB
#    e5 cere prefixul "passage: " la texte și "query: " la întrebări.
# ----------------------------------------------------------------------
def indexeaza():
    embedder = SentenceTransformer(EMB_MODEL)
    client = chromadb.PersistentClient(path=DB_DIR)
    try:
        client.delete_collection(COLLECTION)
    except Exception:
        pass
    col = client.create_collection(COLLECTION, metadata={"hnsw:space": "cosine"})

    docs, metas, ids = [], [], []
    idx = 0
    for cale in glob.glob(os.path.join(TEXTS_DIR, "*.txt")):
        titlu = os.path.splitext(os.path.basename(cale))[0]
        with open(cale, "r", encoding="utf-8") as f:
            text = curata(f.read())
        for nr, frag in enumerate(fragmenteaza(text, CHUNK_WORDS, CHUNK_OVERLAP)):
            docs.append(frag)
            metas.append({"carte": titlu, "fragment": nr})
            ids.append(f"{titlu}_{nr}")
            idx += 1

    if not docs:
        raise SystemExit(f"Niciun fragment de indexat în {TEXTS_DIR}")

    # embed pe loturi ca să nu saturezi memoria
    lot = 64
    for i in range(0, len(docs), lot):
        batch = docs[i:i + lot]
        vec = embedder.encode(
            [f"passage: {d}" for d in batch],
            normalize_embeddings=True,
        ).tolist()
        col.add(
            documents=batch,
            embeddings=vec,
            metadatas=metas[i:i + lot],
            ids=ids[i:i + lot],
        )
        print(f"Indexat {min(i + lot, len(docs))}/{len(docs)}")
    print(f"Gata. {idx} fragmente în {DB_DIR}")


# ----------------------------------------------------------------------
# 3. CĂUTARE + GENERARE RĂSPUNS
# ----------------------------------------------------------------------
def raspunde(intrebare):
    embedder = SentenceTransformer(EMB_MODEL)
    client = chromadb.PersistentClient(path=DB_DIR)
    col = client.get_collection(COLLECTION)

    q_vec = embedder.encode(
        [f"query: {intrebare}"], normalize_embeddings=True
    ).tolist()
    rez = col.query(query_embeddings=q_vec, n_results=TOP_K)

    fragmente = rez["documents"][0]
    surse = rez["metadatas"][0]

    context = "\n\n".join(
        f"[{m['carte']}] {d}" for d, m in zip(fragmente, surse)
    )

    prompt = (
        "Răspunde la întrebare folosind EXCLUSIV textele de mai jos. "
        "Dacă informația nu se află în ele, spune că nu apare în surse. "
        "Citează cartea din care provine fiecare afirmație.\n\n"
        f"### Texte:\n{context}\n\n"
        f"### Întrebare: {intrebare}\n\n"
        "### Răspuns:"
    )

    tokenizer = AutoTokenizer.from_pretrained(GEN_MODEL)
    model = AutoModelForCausalLM.from_pretrained(
        GEN_MODEL, torch_dtype=torch.bfloat16, device_map="auto"
    )
    mesaje = [{"role": "user", "content": prompt}]
    inputs = tokenizer.apply_chat_template(
        mesaje, return_tensors="pt", add_generation_prompt=True
    ).to(model.device)
    with torch.no_grad():
        out = model.generate(inputs, max_new_tokens=400, temperature=0.3, do_sample=True)
    raspuns = tokenizer.decode(out[0][inputs.shape[1]:], skip_special_tokens=True)

    print("\n=== RĂSPUNS ===")
    print(raspuns.strip())
    print("\n=== SURSE FOLOSITE ===")
    for m in surse:
        print(f"- {m['carte']} (fragment {m['fragment']})")


# ----------------------------------------------------------------------
# CLI
# ----------------------------------------------------------------------
if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--index", action="store_true", help="construiește baza de vectori")
    ap.add_argument("--ask", type=str, help="pune o întrebare")
    args = ap.parse_args()

    if args.index:
        indexeaza()
    elif args.ask:
        raspunde(args.ask)
    else:
        print("Folosire: --index (o dată) apoi --ask \"întrebarea ta\"")