fix: revert to last functional version

app.py

@@ -3,9 +3,10 @@ import pickle
 import numpy as np
 import faiss
 import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
 from sentence_transformers import SentenceTransformer
 import gradio as gr
+from threading import Thread
 
 index = faiss.read_index("vector_db/index.faiss")
 with open("vector_db/chunks.pkl", "rb") as f:
@@ -15,7 +16,6 @@ with open("vector_db/metadata.pkl", "rb") as f:
 
 ST = SentenceTransformer("BAAI/bge-large-en-v1.5")
 
-chunk_embeddings = ST.encode(chunks, convert_to_numpy=True)
 model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"
 bnb = BitsAndBytesConfig(
     load_in_4bit=True,
@@ -30,80 +30,76 @@ model = AutoModelForCausalLM.from_pretrained(
     device_map={"": 0},
     torch_dtype=torch.bfloat16
 )
-pipe = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    device_map={"": 0},
-    quantization_config=bnb
-)
 
 SYS = (
-    "### System\n"
-    "You are a legal AI assistant specialized in GDPR/EDPB. "
-    "Answer based only on provided context. If uncertain, say 'I do not know.'\n"
+    "You are a legal AI assistant. Answer the user's question "
+    "based only on the given legal context from GDPR and EDPB documents. "
+    "Be accurate, use clear language, and do not make assumptions. "
+    "If unsure, say: 'I do not know.'"
 )
 
-def retrieve(query: str, k: int = 3):
-    emb = ST.encode([query], convert_to_numpy=True)[0]
+def retrieve(q, k=3):
+    emb = ST.encode(q)
     D, I = index.search(np.array([emb], dtype="float32"), k)
-    idxs = [int(i) for i in I[0]]
-    selected_embs = chunk_embeddings[idxs]
-    sims = np.dot(selected_embs, emb) / (np.linalg.norm(selected_embs, axis=1) * np.linalg.norm(emb))
-    order = np.argsort(-sims)
-    docs, sources = [], []
-    for pos in order:
-        i = idxs[pos]
-        docs.append({"title": f"Chunk {i}", "pages": chunks[i]})
-        sources.append(metadata_dict[i]["source"])
-    return docs, sources
+    
+    docs = []
+    file_sources = []
+    
+    for i in I[0]:
+        chunk = chunks[i]
+        metadata = metadata_dict[i] 
+        docs.append({
+            "title": chunk,
+            "pages": chunk 
+        })
+        file_sources.append(metadata["source"])
+        
+    return docs, file_sources
 
-def make_prompt(question: str, docs: list) -> str:
-    context = "\n---\n".join(f"Title: {d['title']}\n{d['pages']}" for d in docs)
-    return (
-        f"{SYS}"
-        f"### Context\n{context}\n"
-        "### Question\n" + question + "\n"
-        "### Chain of Thought\nThink step by step about relevant legal provisions.\n"
-        "### Answer (JSON)\n{\"answer\": `Your answer`, \"sources\": []}"
-    )
+def make_prompt(q, docs):
+    context = "\n\n".join(f"Title: {doc['title']}\nPages: {doc['pages']}" for doc in docs)
+    return f"{SYS}\n\nContext:\n{context}\n\nQuestion:\n{q}\n\nAnswer:"
 
 @spaces.GPU()
 def qa_fn(question, top_k, temperature, max_tokens):
     docs, file_sources = retrieve(question, top_k)
-    prompt = make_prompt(question, docs)
-    outputs = pipe(
-        prompt,
-        max_new_tokens=max_tokens,
-        do_sample=False,
-        temperature=temperature,
-        top_p=0.9,
-        return_full_text=False
-    )
-    raw = outputs[0]["generated_text"].strip()
-    try:
-        json_start = raw.find('{')
-        json_text = raw[json_start: raw.rfind('}')+1]
-        import json; result = json.loads(json_text.replace('`', '"'))
-        answer = result.get("answer", raw)
-    except Exception:
-        answer = raw
-    return answer, file_sources
+    prompt = make_prompt(question, docs)[:8000]
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}
+    streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True)
+    Thread(target=model.generate, kwargs={
+        **inputs,
+        "streamer": streamer,
+        "max_new_tokens": max_tokens,
+        "do_sample": False,
+        "temperature": temperature,
+        "top_p": 0.9,
+        "eos_token_id": tokenizer.eos_token_id
+    }).start()
+    output = ""
+    for tok in streamer:
+        output += tok
+
+    think_tag_index = output.find("</think>") #change to "Answer:" after testing
+    if think_tag_index != -1:
+        output = output[think_tag_index + len("</think>"):].strip()
+
+    return output, file_sources
+
+outputs_answer = gr.Textbox(label="Answer")
+outputs_sources = gr.JSON(label="Sources (Used Files)")
 
 demo = gr.Interface(
     fn=qa_fn,
     inputs=[
         gr.Textbox(lines=2, label="Your Question"),
-        gr.Slider(1, 15, value=5, step=1, label="Top-K Documents"),
-        gr.Slider(0.1, 1.0, value=0.7, step=0.05, label="Temperature"),
-        gr.Slider(64, 1024, value=512, step=64, label="Max Answer Length")
-    ],
-    outputs=[
-        gr.Textbox(label="Answer"),
-        gr.JSON(label="Sources (Used Files)")
+        gr.Slider(1, 15, value=5, step=1, label="Top-K Documents (How many chunks to include for context)"),
+        gr.Slider(0.1, 1.0, value=0.7, step=0.05, label="Temperature (Higher = more creative, lower = more focused)"),
+        gr.Slider(64, 1024, value=512, step=64, label="Max Answer Length (Maximum tokens to generate)")
     ],
+    outputs=[outputs_answer, outputs_sources],
     title="GDPR Legal Assistant",
-    description="Enhanced RAG with reranking, structured prompts & CoT for precise legal answers.",
+    description="Ask any question about GDPR or EDPB documents. The response includes used files and chunks.",
     allow_flagging="never"
 )