Updated app.py

app.py

@@ -6,27 +6,29 @@ from sklearn.feature_extraction.text import TfidfVectorizer
 import faiss
 import numpy as np
 
-#--------------------MODEL SETUP--------------------
+# ------------------ Model Setup ------------------
+
 MODEL_REPO = "google/gemma-3-1b-it-qat-q4_0-gguf"
 MODEL_PATH = "./gemma-3-1b-it-qat-q4_0/gemma-3-1b-it-q4_0.gguf"
 MODEL_DIR = "./gemma-3-1b-it-qat-q4_0"
+
 DEFAULT_SYSTEM_PROMPT = (
     "You are a Wise Mentor. Speak in a calm and concise manner. "
     "If asked for advice, give a maximum of 3 actionable steps. "
     "Avoid unnecessary elaboration. Decline unethical or harmful requests."
 )
 
-# Huggingface Token and download
+# Hugging Face Token and model download
 hf_token = os.environ.get("HF_TOKEN")
 if not os.path.exists(MODEL_PATH):
     if not hf_token:
-        raise ValueError("HF_TOKEN is missing. Set it as an environment variable")
-    
+        raise ValueError("HF_TOKEN is missing. Set it as an environment variable.")
     login(hf_token)
     snapshot_download(repo_id=MODEL_REPO, local_dir=MODEL_DIR, local_dir_use_symlinks=False)
 
-#--------------------RAG SETUP------------------------
-documents = [] # stores all chat turns
+# ------------------ RAG Setup ------------------
+
+documents = []  # stores (user, bot) tuples
 vectorizer = TfidfVectorizer()
 index = None
 
@@ -34,40 +36,41 @@ def update_rag_index():
     global index
     if not documents:
         return
-    vectors = vectorizer.fit_transform(documents).toarray().astype('float32')
+    flat_docs = [f"user: {u} bot: {b}" for u, b in documents]
+    vectors = vectorizer.fit_transform(flat_docs).toarray().astype('float32')
     index = faiss.IndexFlatL2(vectors.shape[1])
     index.add(vectors)
 
-def retrive_relvant_docs(query, k=2):
+def retrieve_relevant_docs(query, k=3):
     if not documents or index is None:
-        return ""
-    
-    query_vac = vectorizer.transform([query]).toarray().astype('float32')
-    D, I = index.search(query_vac, k)
-    return "\n".join(documents[i] for i in I[0] if i < len(documents))
+        return []
+    flat_docs = [f"user: {u} bot: {b}" for u, b in documents]
+    query_vec = vectorizer.transform([query]).toarray().astype('float32')
+    D, I = index.search(query_vec, k)
+    return [documents[i] for i in I[0] if i < len(documents)]
 
+# ------------------ Context Estimation ------------------
 
-#-----------------------CONTEXT LENGTH ESTIMATION---------------------
-def estimate_n_ctx(full_prompt, buffer = 500):
+def estimate_n_ctx(full_prompt, buffer=500):
     total_tokens = len(full_prompt.split())
-    return min(3500, total_tokens+buffer)
+    return min(3500, total_tokens + buffer)
+
+# ------------------ Chat Function ------------------
 
-#-----------------------CHAT FUNCTION-----------------------
 def chat(user_input, history, system_prompt):
-    relevent_context = retrive_relvant_docs(user_input)
-    formatted_turns = "".join([f"<user>{u}</user><bot>{b}</bot>" for u, b in relevent_context])
+    relevant_context = retrieve_relevant_docs(user_input)
+    formatted_turns = "".join([f"<user>{u}</user><bot>{b}</bot>" for u, b in relevant_context])
 
     full_prompt = (
-        f"<s>[INST] <<SYS>>\n{system_prompt}\nContext:\n{relevent_context}\n<</SYS>>\n"
+        f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n"
         f"{formatted_turns}<user>{user_input}[/INST]"
     )
 
-    # Dynamic estimate n_ctx
-    n_ctx = estimate_n_ctx(full_prompt=full_prompt)
+    n_ctx = estimate_n_ctx(full_prompt)
 
     llm = Llama(
-        model_path= MODEL_PATH,
-        n_ctx = n_ctx,
+        model_path=MODEL_PATH,
+        n_ctx=n_ctx,
         n_threads=2,
         n_batch=128
     )
@@ -75,22 +78,22 @@ def chat(user_input, history, system_prompt):
     output = llm(full_prompt, max_tokens=256, stop=["</s>", "<user>"])
     bot_reply = output["choices"][0]["text"].strip()
 
-    documents.append(f"user: {user_input} bot: {bot_reply}")
+    documents.append((user_input, bot_reply))
     update_rag_index()
 
     history.append((user_input, bot_reply))
     return "", history
 
-#-----------------------UI---------------------
+# ------------------ UI ------------------
+
 with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 Persona Agent with Mini-RAG + Dynamic Context")
-    with gr.Row():
-        system_prompt_box = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=3)
+    gr.Markdown("## 🤖 Prompt-Engineered Persona Agent with Mini-RAG")
+    system_prompt_box = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=3)
     chatbot = gr.Chatbot()
     msg = gr.Textbox(label="Your Message")
-    clear = gr.Button("🗑️ Clear")
+    clear = gr.Button("🗑️ Clear Chat")
 
     msg.submit(chat, [msg, chatbot, system_prompt_box], [msg, chatbot])
     clear.click(lambda: [], None, chatbot)
 
-demo.launch()
+demo.launch()