Update app.py

app.py

@@ -1,70 +1,158 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
-
-
-def respond(
-    message,
-    history: list[dict[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-    hf_token: gr.OAuthToken,
-):
-    """
-    For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-    """
-    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")
-
-    messages = [{"role": "system", "content": system_message}]
-
-    messages.extend(history)
-
-    messages.append({"role": "user", "content": message})
-
-    response = ""
-
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        choices = message.choices
-        token = ""
-        if len(choices) and choices[0].delta.content:
-            token = choices[0].delta.content
-
-        response += token
-        yield response
-
-
-"""
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-chatbot = gr.ChatInterface(
-    respond,
-    type="messages",
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
-)
+import fitz  # PyMuPDF
+import re
+import numpy as np
+import faiss
+import os
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+from huggingface_hub import login
 
+# -----------------------------
+# PDF Text Loader
+# -----------------------------
+def load_pdf_text(file_obj):
+    doc = fitz.open(stream=file_obj.read(), filetype="pdf")
+    text = ""
+    for page in doc:
+        text += page.get_text()
+    if not text.strip():
+        raise ValueError("No text found in PDF.")
+    return text
+
+# -----------------------------
+# Chunk Text
+# -----------------------------
+def chunk_text(text, max_tokens=200):
+    sentences = re.split(r'(?<=[.!?]) +', text)
+    chunks, current_chunk = [], []
+    current_len = 0
+    for sentence in sentences:
+        word_count = len(sentence.split())
+        if current_len + word_count > max_tokens:
+            chunks.append(" ".join(current_chunk))
+            current_chunk = [sentence]
+            current_len = word_count
+        else:
+            current_chunk.append(sentence)
+            current_len += word_count
+    if current_chunk:
+        chunks.append(" ".join(current_chunk))
+    return chunks
+
+# -----------------------------
+# Simple Vector Store
+# -----------------------------
+class SimpleVectorStore:
+    def __init__(self, dim):
+        self.dim = dim
+        self.vectors = []
+        self.metadata = []
+        self.index = None
+
+    def add(self, vectors, metas):
+        for v, m in zip(vectors, metas):
+            vec = np.array(v, dtype=np.float32)
+            self.vectors.append(vec)
+            self.metadata.append(m)
+        if self.vectors:
+            self.index = faiss.IndexFlatL2(self.dim)
+            self.index.add(np.stack(self.vectors))
+
+    def search(self, query_vector, k=5):
+        query_vector = np.array(query_vector, dtype=np.float32).reshape(1, -1)
+        D, I = self.index.search(query_vector, k)
+        results = [self.metadata[i] for i in I[0]]
+        return results
+
+# -----------------------------
+# Index PDF
+# -----------------------------
+def index_pdf(file_obj):
+    text = load_pdf_text(file_obj)
+    chunks = chunk_text(text)
+    embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+    vectors = embed_model.encode(chunks)
+    store = SimpleVectorStore(dim=vectors.shape[1])
+    store.add(vectors, chunks)
+    return embed_model, store
+
+# -----------------------------
+# Load LLaMA Model
+# -----------------------------
+def load_llm():
+    model_id = "meta-llama/Llama-3.2-3b-instruct"
+    hf_token = os.getenv("HF_TOKEN")
+    if not hf_token:
+        raise ValueError("HF_TOKEN is not set. Please add it in Hugging Face Secrets.")
+    login(hf_token)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+    llm = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        token=hf_token
+    )
+    return tokenizer, llm
+
+# -----------------------------
+# HyDE + Answer Query
+# -----------------------------
+def answer_query(file_obj, question):
+    try:
+        embed_model, store = index_pdf(file_obj)
+        tokenizer, llm = load_llm()
+
+        # ---- Step 1: HyDE hypothetical answer ----
+        hyde_prompt = f"""
+        [INST] Write a detailed hypothetical answer to this question:
+        {question}
+        Answer: [/INST]
+        """
+        inputs = tokenizer(hyde_prompt, return_tensors="pt").to(llm.device)
+        hyde_out = llm.generate(**inputs, max_new_tokens=200)
+        hypo_answer = tokenizer.decode(hyde_out[0], skip_special_tokens=True)
+
+        # ---- Step 2: Embed hypothetical answer ----
+        query_vec = embed_model.encode([hypo_answer])[0]
+
+        # ---- Step 3: Retrieve top chunks ----
+        relevant_chunks = store.search(query_vec, k=5)
+        context = "\n".join(relevant_chunks)
+
+        # ---- Step 4: Final Answer ----
+        final_prompt = f"""
+        [INST] You are a helpful tutor. Based only on the context below, answer the question.
+        If context does not have the info, say "I could not find this in the text."
+        Context:
+        {context}
+        Question: {question}
+        Answer: [/INST]
+        """
+        inputs = tokenizer(final_prompt, return_tensors="pt", truncation=True).to(llm.device)
+        outputs = llm.generate(**inputs, max_new_tokens=300, temperature=0.7, top_p=0.9, do_sample=True)
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+        if "Answer:" in answer:
+            answer = answer.split("Answer:")[-1].strip()
+
+        return answer
+
+    except Exception as e:
+        return f"⚠️ Error: {e}"
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
 with gr.Blocks() as demo:
-    with gr.Sidebar():
-        gr.LoginButton()
-    chatbot.render()
+    gr.Markdown("## 📚 HyDE RAG Chatbot (PDF Tutor)")
+    file_input = gr.File(label="Upload PDF", type="filepath")
+    question = gr.Textbox(label="Ask a Question")
+    answer = gr.Textbox(label="Answer", interactive=False)
+    btn = gr.Button("Get Answer")
 
+    btn.click(fn=answer_query, inputs=[file_input, question], outputs=answer)
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()