Update app.py

app.py

@@ -1,58 +1,57 @@
 import gradio as gr
 import os
-import faiss
-import torch
 import json
+import faiss
 import numpy as np
-from huggingface_hub import hf_hub_download, InferenceClient
+import torch
 from sentence_transformers import SentenceTransformer
+from huggingface_hub import InferenceClient, hf_hub_download
 
-# Hugging Face Credentials
+# 🔹 Hugging Face Credentials
 HF_REPO = "Futuresony/future_ai_12_10_2024.gguf"
-HF_FAISS_FILE = "asa_faiss.index"
-api_token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
+HF_TOKEN = os.getenv('HUGGINGFACEHUB_API_TOKEN')  # Store your token as an environment variable for security
 
-# Load the FAISS index from Hugging Face
-faiss_local_path = hf_hub_download(repo_id=HF_REPO, filename=HF_FAISS_FILE, repo_type="model", token=api_token)
-index = faiss.read_index(faiss_local_path)
+# 🔹 FAISS Index Path
+FAISS_PATH = "asa_faiss.index"
 
-# Load the same embedding model used for FAISS
-embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+# 🔹 Load Sentence Transformer for Embeddings
+embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# Hugging Face model client
-client = InferenceClient(model=HF_REPO, token=api_token)
+# 🔹 Load FAISS Index from Hugging Face
+faiss_local_path = hf_hub_download(HF_REPO, "asa_faiss.index", token=HF_TOKEN)
+faiss_index = faiss.read_index(faiss_local_path)
 
-def retrieve_relevant_context(query, top_k=3):
-    """Retrieve the most relevant text chunks from FAISS."""
-    query_embedding = embedding_model.encode([query]).astype(np.float32)
-    distances, indices = index.search(query_embedding, top_k)
+# 🔹 Initialize Hugging Face Model Client
+client = InferenceClient(model=HF_REPO, token=HF_TOKEN)
+
+# 🔹 Retrieve Relevant FAISS Context
+def retrieve_relevant_context(user_query, top_k=3):
+    query_embedding = embedder.encode([user_query], convert_to_tensor=True).cpu().numpy()
+    distances, indices = faiss_index.search(query_embedding, top_k)
 
     retrieved_texts = []
-    for idx in indices[0]:  # Get the closest matches
-        if idx != -1:  # Valid match
-            retrieved_texts.append(f"Relevant info: {idx}")
+    for idx in indices[0]:  # Extract top_k results
+        if idx != -1:  # Ensure valid index
+            retrieved_texts.append(f"Example: {idx} → {idx}")  # Customize how retrieved data appears
 
-    return "\n".join(retrieved_texts) if retrieved_texts else "No relevant info found."
+    return "\n".join(retrieved_texts) if retrieved_texts else "No relevant data found."
 
-def format_alpaca_prompt(user_input, system_prompt, history):
-    """Formats input in Alpaca/LLaMA style"""
-    history_str = "\n".join([f"### Instruction:\n{h[0]}\n### Response:\n{h[1]}" for h in history])
-    prompt = f"""{system_prompt}
-{history_str}
+# 🔹 Format Model Prompt with FAISS Guidance
+def format_prompt(user_input, system_prompt, history):
+    retrieved_context = retrieve_relevant_context(user_input)
 
-### Instruction:
-{user_input}
+    faiss_instruction = (
+        "Use the following example responses as a guide for formatting and writing style:\n"
+        f"{retrieved_context}\n\n"
+        "### Instruction:\n"
+        f"{user_input}\n\n### Response:"
+    )
 
-### Response:
-"""
-    return prompt
+    return faiss_instruction
 
+# 🔹 Chatbot Response Function
 def respond(message, history, system_message, max_tokens, temperature, top_p):
-    # 🔹 Retrieve relevant info from FAISS
-    retrieved_context = retrieve_relevant_context(message)
-    
-    # 🔹 Include retrieved context in the prompt
-    full_prompt = f"{retrieved_context}\n\n{format_alpaca_prompt(message, system_message, history)}"
+    full_prompt = format_prompt(message, system_message, history)
 
     response = client.text_generation(
         full_prompt,
@@ -61,17 +60,18 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         top_p=top_p,
     )
 
-    # ✅ Extract only the response
+    # ✅ Extract only model-generated response
     cleaned_response = response.split("### Response:")[-1].strip()
     
-    history.append((message, cleaned_response))  # ✅ Update history with the new message and response
+    history.append((message, cleaned_response))  # ✅ Update chat history
     
-    yield cleaned_response  # ✅ Output only the answer
+    yield cleaned_response  # ✅ Output the response
 
+# 🔹 Gradio Chat Interface
 demo = gr.ChatInterface(
     respond,
     additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
+        gr.Textbox(value="You are a helpful AI trained to follow FAISS-based writing styles.", label="System message"),
         gr.Slider(minimum=1, maximum=250, value=128, step=1, label="Max new tokens"),
         gr.Slider(minimum=0.1, maximum=4.0, value=0.9, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.99, step=0.01, label="Top-p (nucleus sampling)"),