krishankula/rag_base_model

1import os2import torch3import faiss4import pickle5import gradio as gr6from transformers import AutoTokenizer, AutoModelForCausalLM7from sentence_transformers import SentenceTransformer8 9# ==============================10# 1️⃣ DEVICE SETUP11# ==============================12device = "cuda" if torch.cuda.is_available() else "cpu"13 14# ==============================15# 2️⃣ LOAD BASE MODEL (NO accelerate needed)16# ==============================17model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"18 19tokenizer = AutoTokenizer.from_pretrained(model_name)20 21model = AutoModelForCausalLM.from_pretrained(22    model_name,23    dtype=torch.float16 if device == "cuda" else torch.float3224)25 26model.to(device)27model.eval()28 29# ==============================30# 3️⃣ LOAD EMBEDDING MODEL31# ==============================32embed_model = SentenceTransformer(33    "sentence-transformers/all-MiniLM-L6-v2",34    device=device35)36 37# ==============================38# 4️⃣ LOAD FAISS + DOCUMENTS SAFELY39# ==============================40if not os.path.exists("faiss_index.bin") or not os.path.exists("documents.pkl"):41    raise FileNotFoundError(42        "faiss_index.bin or documents.pkl not found. "43        "Upload them to your HF Space repo."44    )45 46index = faiss.read_index("faiss_index.bin")47 48with open("documents.pkl", "rb") as f:49    documents = pickle.load(f)50 51# ==============================52# 5️⃣ RETRIEVAL FUNCTION53# ==============================54def retrieve_context(query, top_k=2):55    query_embedding = embed_model.encode(56        [query],57        convert_to_numpy=True58    )59 60    distances, indices = index.search(query_embedding, top_k)61    retrieved_docs = [documents[i] for i in indices[0]]62 63    return "\n\n".join(retrieved_docs)64 65# ==============================66# 6️⃣ RAG GENERATION67# ==============================68def rag_answer(question):69    context = retrieve_context(question)70 71    prompt = f"""<|system|>72You are a legal assistant.73Answer ONLY using the provided context.74If the answer is not in the context, say:75"I cannot find evidence in the retrieved documents."76</s>77<|user|>78Context:79{context}80 81Question:82{question}83</s>84<|assistant|>85"""86 87    inputs = tokenizer(prompt, return_tensors="pt").to(device)88 89    with torch.no_grad():90        outputs = model.generate(91            **inputs,92            max_new_tokens=150,93            temperature=0.2,94            repetition_penalty=1.2,95            do_sample=True,96            eos_token_id=tokenizer.eos_token_id,97            pad_token_id=tokenizer.eos_token_id98        )99 100    decoded = tokenizer.decode(outputs[0], skip_special_tokens=True)101 102    # Extract only assistant answer safely103    if "<|assistant|>" in decoded:104        decoded = decoded.split("<|assistant|>")[-1]105 106    return decoded.strip()107 108# ==============================109# 7️⃣ GRADIO UI110# ==============================111demo = gr.Interface(112    fn=rag_answer,113    inputs=gr.Textbox(114        lines=3,115        label="Legal Question",116        placeholder="e.g., What are the penalties under BIPA?"117    ),118    outputs=gr.Textbox(label="Grounded Answer"),119    title="⚖️ TinyLlama Legal Hallucination-Reduced(Base+Rag) Model",120    description="RAG system using FAISS retrieval over US legal documents."121)122 123demo.launch()