Update app.py

app.py

@@ -3,14 +3,14 @@ import gradio as gr
 import numpy as np
 import faiss
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 # =========================================================
 # CONFIG
 # =========================================================
 EMBEDDING_MODEL = "all-MiniLM-L6-v2"
-GENERATION_MODEL = "google/flan-t5-base"
+GENERATION_MODEL = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 TOP_K = 3
 MAX_NEW_TOKENS = 200
 
@@ -67,7 +67,7 @@ def chunk_documents(documents, max_words=100):
     return chunks
 
 # =========================================================
-# BUILD VECTOR STORE
+# VECTOR STORE
 # =========================================================
 def build_vector_store(chunks):
     embedder = SentenceTransformer(EMBEDDING_MODEL)
@@ -88,51 +88,45 @@ def retrieve(query, index, chunks, embedder):
     return [chunks[i] for i in idxs[0]]
 
 # =========================================================
-# GENERATION ENGINE (SEQ2SEQ – CORRECT)
+# GENERATION ENGINE (LLAMA – CAUSAL)
 # =========================================================
 class AnswerGenerator:
     def __init__(self):
         self.tokenizer = AutoTokenizer.from_pretrained(GENERATION_MODEL)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(GENERATION_MODEL)
+        self.model = AutoModelForCausalLM.from_pretrained(GENERATION_MODEL)
         self.device = "cuda" if torch.cuda.is_available() else "cpu"
         self.model.to(self.device)
 
     def generate(self, query, context_chunks):
         context = "\n".join([c["text"] for c in context_chunks[:3]])
 
-        prompt = f"""
-Based on the information below, answer clearly and professionally.
-
-Information:
-{context}
-
-Question: {query}
-
-Answer:
-"""
-
-        inputs = self.tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=512
+        prompt = (
+            "<|system|>\n"
+            "You are an AI professional profile assistant.\n"
+            "Answer using ONLY the information provided.\n"
+            "Be concise and recruiter-facing.\n"
+            "<|user|>\n"
+            f"Information:\n{context}\n\n"
+            f"Question: {query}\n"
+            "<|assistant|>\n"
         )
+
+        inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=1024)
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
 
         with torch.no_grad():
             outputs = self.model.generate(
                 **inputs,
                 max_new_tokens=MAX_NEW_TOKENS,
-                num_beams=4,
-                temperature=0.7,
-                do_sample=True,
-                top_p=0.9,
-                no_repeat_ngram_size=3,
-                early_stopping=True
+                do_sample=False,
+                temperature=0.2,
+                repetition_penalty=1.1,
+                eos_token_id=self.tokenizer.eos_token_id
             )
 
-        answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return answer.strip()
+        decoded = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        answer = decoded.split("<|assistant|>")[-1].strip()
+        return answer
 
 # =========================================================
 # QUERY ENHANCEMENT
@@ -152,7 +146,7 @@ def enhance_query(query):
     return query
 
 # =========================================================
-# INITIALIZE SYSTEM (ONCE)
+# INITIALIZE SYSTEM
 # =========================================================
 documents = load_documents()
 chunks = chunk_documents(documents)
@@ -165,7 +159,7 @@ def answer_question(user_query):
     return generator.generate(enhanced, retrieved)
 
 # =========================================================
-# CHATBOT RESPONSE FUNCTION
+# CHATBOT RESPONSE
 # =========================================================
 def respond(message, history):
     if not message or not message.strip():
@@ -176,16 +170,14 @@ def respond(message, history):
     return history, history
 
 # =========================================================
-# CHATBOT UI (GRADIO TEMPLATE COMPATIBLE)
+# CHATBOT UI
 # =========================================================
 with gr.Blocks(title="AI Twin Chatbot – Aniket Sirsikar") as demo:
 
     gr.Markdown(
         """
-        # 🤖 AI Twin – Chatbot Mode
+        # 🤖 AI Twin – Chatbot (LLaMA)
         **Product Manager | Data Analytics | AI-Powered Systems**
-
-        Ask me about my skills, projects, experience, education, or career goals.
         """
     )
 
@@ -199,16 +191,7 @@ with gr.Blocks(title="AI Twin Chatbot – Aniket Sirsikar") as demo:
 
     send = gr.Button("Send")
 
-    send.click(
-        fn=respond,
-        inputs=[user_input, state],
-        outputs=[chatbot, state]
-    )
-
-    user_input.submit(
-        fn=respond,
-        inputs=[user_input, state],
-        outputs=[chatbot, state]
-    )
+    send.click(respond, [user_input, state], [chatbot, state])
+    user_input.submit(respond, [user_input, state], [chatbot, state])
 
 demo.queue().launch()