Update app.py

app.py

@@ -1,17 +1,18 @@
-import gradio as gr
 import torch
 from transformers import AutoTokenizer, pipeline
 from sentence_transformers import SentenceTransformer
 import faiss
 import numpy as np
+import gradio as gr
+from typing import List
 
 # Configuration
 class Config:
     model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
     embedding_model = "all-MiniLM-L6-v2"
-    vector_dim = 384 
-    top_k = 3 
-    chunk_size = 256 
+    vector_dim = 384  # Sentence Transformer embedding dimension
+    top_k = 3  # Retrieve top 3 relevant chunks
+    chunk_size = 256  # Text chunk size
 
 # Vector Database
 class VectorDB:
@@ -27,7 +28,7 @@ class VectorDB:
         self.index.add(embedding)
         self.texts.append(text)
     
-    def search(self, query: str):
+    def search(self, query: str) -> List[str]:
         if self.index.ntotal == 0:
             return []
         query_embedding = self.embedding_model.encode([query])[0]
@@ -42,7 +43,7 @@ class TinyChatModel:
         self.tokenizer = AutoTokenizer.from_pretrained(Config.model_name)
         self.pipe = pipeline("text-generation", model=Config.model_name, torch_dtype=torch.bfloat16, device_map="auto")
 
-    def generate_response(self, message: str, context: str = ""):
+    def generate_response(self, message: str, context: str = "") -> str:
         messages = [{"role": "user", "content": message}]
         if context:
             messages.insert(0, {"role": "system", "content": f"Context:\n{context}"})
@@ -54,32 +55,35 @@ class TinyChatModel:
 vector_db = VectorDB()
 chat_model = TinyChatModel()
 
-def chat_interface(user_input):
-    context = "\n".join(vector_db.search(user_input))
-    response = chat_model.generate_response(user_input, context)
+# Function to handle context addition and chat
+def chat_function(user_input: str, context: str = ""):
+    if context:
+        vector_db.add_text(context)
+    
+    # Search relevant context
+    context_text = "\n".join(vector_db.search(user_input))
+    response = chat_model.generate_response(user_input, context_text)
     vector_db.add_text(f"User: {user_input}\nAssistant: {response}")
+    
     return response
 
-def add_text_interface(text):
-    vector_db.add_text(text)
-    return "Text added to memory!"
+# Gradio Interface
+def gradio_interface(user_input: str, context: str = ""):
+    response = chat_function(user_input, context)
+    return response
 
-# Gradio UI
-demo = gr.Blocks()
-with demo:
-    gr.Markdown("# 🦙 TinyChat - AI Chatbot")
-    with gr.Row():
-        chatbot = gr.Chatbot()
+# Create Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("# TinyChat: A Conversational AI")
     with gr.Row():
-        user_input = gr.Textbox(label="Your Message")
-        send_btn = gr.Button("Send")
-    with gr.Row():
-        add_text_input = gr.Textbox(label="Add Knowledge to AI")
-        add_text_btn = gr.Button("Add Text")
-    
-    send_btn.click(chat_interface, inputs=user_input, outputs=chatbot)
-    add_text_btn.click(add_text_interface, inputs=add_text_input, outputs=gr.Textbox())
+        with gr.Column():
+            user_input = gr.Textbox(label="User Input", placeholder="Ask anything...")
+            context_input = gr.Textbox(label="Optional Context", placeholder="Paste context here (optional)", lines=3)
+            submit_button = gr.Button("Send")
+            output = gr.Textbox(label="Response", placeholder="Assistant's reply will appear here...")
+            
+            submit_button.click(fn=gradio_interface, inputs=[user_input, context_input], outputs=output)
 
-# Launch
+# Run the Gradio app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=True)