feat: restart and build from ground on

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

app.py

@@ -1,113 +1,75 @@
-# -*- coding: utf-8 -*-
 import spaces
-import os, logging, traceback, pickle, gc
-import gradio as gr
-import torch
-import faiss
+import pickle
 import numpy as np
-from transformers import AutoTokenizer, AutoModelForCausalLM
+import faiss
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TextIteratorStreamer
 from sentence_transformers import SentenceTransformer
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-
-VECTOR_DB_DIR      = "vector_db"
-EMBEDDING_MODEL    = "BAAI/bge-large-en-v1.5"
-GEN_QA_MODEL       = "Qwen/Qwen2-7B-Instruct"
-MODEL_CONTEXT_SIZE = 10000
-CHUNK_SIZE         = 512
-PROMPT_RESERVE     = 1024
+import gradio as gr
+from threading import Thread
 
-index = None
-document_metadata = []
-all_chunks = []
-emb_model = None
-gen_tokenizer = None
-gen_model = None
+index = faiss.read_index("vector_db/index.faiss")
+with open("vector_db/chunks.pkl", "rb") as f:
+    chunks = pickle.load(f)
 
-splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+ST = SentenceTransformer("BAAI/bge-large-en-v1.5")
 
-def calculate_top_k():
-    return (MODEL_CONTEXT_SIZE - PROMPT_RESERVE) // CHUNK_SIZE
+model_id = "Qwen/Qwen2.5-7B-Instruct"
+bnb = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    quantization_config=bnb,
+    device_map={"": 0},
+    torch_dtype=torch.bfloat16
+)
 
-def initialize_models():
-    global emb_model, gen_tokenizer, gen_model
-    emb_model = SentenceTransformer(EMBEDDING_MODEL, device="cpu")
-    gen_tokenizer = AutoTokenizer.from_pretrained(
-        GEN_QA_MODEL,
-        use_fast=False,
-        trust_remote_code=True
-    )
-    gen_tokenizer.pad_token = gen_tokenizer.eos_token
-    gen_model = AutoModelForCausalLM.from_pretrained(
-        GEN_QA_MODEL,
-        trust_remote_code=True,
-        device_map="cpu",
-        torch_dtype=torch.float16,
-        load_in_4bit=True,
-        low_cpu_mem_usage=True
-    )
-    gen_model.eval()
-    return "Models loaded."
+SYS = "You are a specialized assistant for answering questions related to legal texts from the GDPR (General Data Protection Regulation) and several Documents of the EDPB (European Data Protection Board). " \
+"Your task is to provide precise and detailed answers based on the provided excerpts from the documents. " \
+"Ensure that you clearly and understandably explain the relevant legal concepts. If you do not know the answer or if the information is insufficient, respond with: 'I do not know.' " \
+"Avoid giving inaccurate or speculative answers."
 
-@spaces.GPU(duration=120)
-def load_faiss_database(progress=gr.Progress()):
-    global index, document_metadata, all_chunks
-    progress(0, "Reading FAISS index...")
-    idx_path = os.path.join(VECTOR_DB_DIR, "index.faiss")
-    if not os.path.exists(idx_path):
-        idx_path = os.path.join(VECTOR_DB_DIR, "faiss_index.idx")
-    cpu_index = faiss.read_index(idx_path)
-    progress(30, "Moving FAISS index to GPU...")
-    res = faiss.StandardGpuResources()
-    index = faiss.index_cpu_to_gpu(res, 0, cpu_index)
-    progress(60, "Loading chunks & metadata...")
-    with open(os.path.join(VECTOR_DB_DIR, "chunks.pkl"), "rb") as f:
-        all_chunks = pickle.load(f)
-    with open(os.path.join(VECTOR_DB_DIR, "metadata.pkl"), "rb") as f:
-        document_metadata = pickle.load(f)
-    progress(100, "FAISS DB ready.")
-    return f"FAISS DB: {len(all_chunks)} chunks."
+def retrieve(q, k=3):
+    emb = ST.encode(q)
+    D, I = index.search(np.array([emb], dtype="float32"), k)
+    return [chunks[i] for i in I[0]]
 
-@spaces.GPU(duration=120)
-def generate_answer(question, db_loaded):
-    global emb_model, gen_model
-    if not db_loaded:
-        return "Please initialize FAISS DB first."
-    emb_model.to("cuda")
-    gen_model.to("cuda")
-    torch.cuda.empty_cache()
-    gc.collect()
-    q_emb = emb_model.encode([question], convert_to_numpy=True)
-    dists, ids = index.search(q_emb.astype(np.float32), calculate_top_k())
-    ctx, sources = [], set()
-    for i in ids[0]:
-        m = document_metadata[i]
-        info = f"{m['source']} (p{m['page']})"
-        ctx.append(f"{info}: {all_chunks[i]}")
-        sources.add(info)
-    docs = splitter.split_text("\n\n".join(ctx))
-    full_context = "\n\n".join(docs)
-    messages = [
-        {"role":"system","content":"You are a GDPR/EDPB expert."},
-        {"role":"user","content":f"Context:\n{full_context}\n\nQ: {question}"}
-    ]
-    prompt = gen_tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
-    inputs = gen_tokenizer(prompt, return_tensors="pt", padding=True).to("cuda")
-    out = gen_model.generate(**inputs, max_new_tokens=PROMPT_RESERVE, do_sample=True)
-    text = gen_tokenizer.decode(out[0], skip_special_tokens=True).split("Assistant:")[-1].strip()
-    return f"Answer:\n{text}\n\nSources:\n- " + "\n- ".join(sources)
+def make_prompt(q, docs):
+    return SYS + "\n\nContext:\n" + "\n".join(docs) + f"\n\nQuestion: {q}\nAnswer:"
 
-with gr.Blocks(theme=gr.themes.Soft(), title="GDPR/EDPB Assistant") as demo:
-    status      = gr.Textbox(label="Status", interactive=False)
-    init_btn    = gr.Button("Initialize FAISS DB")
-    db_loaded   = gr.State(False)
-    question    = gr.Textbox(label="Legal Question", lines=3)
-    submit_btn  = gr.Button("Submit Question")
-    answer      = gr.Textbox(label="Answer", lines=12, interactive=False)
+@spaces.GPU
+def qa_fn(question: str) -> str:
+    docs = retrieve(question, 10)
+    prompt = make_prompt(question, docs)[:8000]
+    inputs = tokenizer(prompt, return_tensors="pt")
+    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": 512,
+        "do_sample": True,
+        "temperature": 0.7,
+        "top_p": 0.9,
+        "eos_token_id": tokenizer.eos_token_id
+    }).start()
+    out = ""
+    for tok in streamer:
+        out += tok
+    return out
 
-    demo.load(initialize_models, outputs=status)
-    init_btn.click(load_faiss_database, outputs=status)
-    init_btn.click(lambda: True, outputs=db_loaded)
-    submit_btn.click(generate_answer, inputs=[question, db_loaded], outputs=answer)
+demo = gr.Interface(
+    fn=qa_fn,
+    inputs=gr.Textbox(lines=2, label="Your question"),
+    outputs=gr.Textbox(lines=10, label="Answer"),
+    title="GDPR QA (RAG)",
+    description="Ask questions on GDPR; answers are grounded in EDPB document chunks."
+)
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)

requirements.txt

@@ -1,10 +1,9 @@
-spaces
 torch
 transformers
 sentence-transformers
-langchain
 faiss-gpu
 gradio
 numpy<2
-accelerate
 bitsandbytes
+accelerate
+spaces