Upload 2 files

app.py

@@ -1,15 +1,14 @@
 import os
+import json
+import numpy as np
 import gradio as gr
-import chromadb
 from huggingface_hub import InferenceClient
 
 # --- Configuration ---
 KNOWLEDGE_BASE_DIR = "knowledge_base"
-COLLECTION_NAME = "ai_twin_kb"
 
 # --- Step 1: Load documents ---
 def load_documents():
-    """Loads all .txt files from the knowledge base directory."""
     documents = []
     filenames = []
     for filename in os.listdir(KNOWLEDGE_BASE_DIR):
@@ -24,7 +23,6 @@ def load_documents():
 
 # --- Step 2: Chunk documents ---
 def chunk_text(text, chunk_size=500, overlap=100):
-    """Splits text into overlapping chunks."""
     chunks = []
     start = 0
     while start < len(text):
@@ -33,77 +31,67 @@ def chunk_text(text, chunk_size=500, overlap=100):
         start += chunk_size - overlap
     return chunks
 
-# --- Step 3: Get embeddings via HF API (no local model!) ---
+# --- Step 3: Get embeddings via HF API ---
 def get_embeddings(texts, client):
-    """Gets embeddings from Hugging Face Inference API."""
-    embeddings = client.feature_extraction(
-        texts, 
-        model="sentence-transformers/all-MiniLM-L6-v2"
-    )
-    # The API returns nested lists, convert to list of lists
-    result = []
-    for emb in embeddings:
-        if isinstance(emb[0], list):
-            # Mean pooling if token-level embeddings returned
-            import numpy as np
-            arr = np.array(emb)
-            pooled = arr.mean(axis=0).tolist()
-            result.append(pooled)
-        else:
-            result.append(emb)
-    return result
-
-# --- Step 4: Build vector store ---
-def build_vector_store(documents, filenames, client):
-    """Creates embeddings via API and stores them in ChromaDB."""
-    chroma_client = chromadb.Client()
-    try:
-        chroma_client.delete_collection(COLLECTION_NAME)
-    except:
-        pass
-    collection = chroma_client.create_collection(name=COLLECTION_NAME)
+    embeddings = []
+    for text in texts:
+        response = client.feature_extraction(text, model="sentence-transformers/all-MiniLM-L6-v2")
+        emb = np.array(response)
+        if emb.ndim == 2:
+            emb = emb.mean(axis=0)
+        embeddings.append(emb)
+    return np.array(embeddings)
+
+# --- Step 4: Simple vector search with numpy ---
+def cosine_similarity(a, b):
+    a_norm = a / (np.linalg.norm(a, axis=-1, keepdims=True) + 1e-10)
+    b_norm = b / (np.linalg.norm(b, axis=-1, keepdims=True) + 1e-10)
+    return np.dot(a_norm, b_norm.T)
+
+class SimpleVectorStore:
+    def __init__(self):
+        self.chunks = []
+        self.sources = []
+        self.embeddings = None
+    
+    def add(self, chunks, sources, embeddings):
+        self.chunks = chunks
+        self.sources = sources
+        self.embeddings = embeddings
     
+    def search(self, query_embedding, top_k=3):
+        scores = cosine_similarity(query_embedding.reshape(1, -1), self.embeddings)[0]
+        top_indices = np.argsort(scores)[-top_k:][::-1]
+        results = [(self.chunks[i], self.sources[i], float(scores[i])) for i in top_indices]
+        return results
+
+# --- Step 5: Build the knowledge store ---
+def build_store(documents, filenames, client):
     all_chunks = []
-    all_ids = []
-    all_metadata = []
-    chunk_id = 0
+    all_sources = []
     
     for doc, fname in zip(documents, filenames):
         chunks = chunk_text(doc)
         for chunk in chunks:
             all_chunks.append(chunk)
-            all_ids.append(f"chunk_{chunk_id}")
-            all_metadata.append({"source": fname})
-            chunk_id += 1
+            all_sources.append(fname)
     
-    print(f"Generating embeddings for {len(all_chunks)} chunks via API...")
-    # Process in batches to avoid API limits
-    batch_size = 16
-    all_embeddings = []
-    for i in range(0, len(all_chunks), batch_size):
-        batch = all_chunks[i:i+batch_size]
-        batch_embeddings = get_embeddings(batch, client)
-        all_embeddings.extend(batch_embeddings)
-        print(f"  Processed {min(i+batch_size, len(all_chunks))}/{len(all_chunks)} chunks")
+    print(f"Embedding {len(all_chunks)} chunks via API...")
+    embeddings = get_embeddings(all_chunks, client)
+    print("Embeddings complete.")
     
-    collection.add(
-        documents=all_chunks,
-        embeddings=all_embeddings,
-        ids=all_ids,
-        metadatas=all_metadata
-    )
+    store = SimpleVectorStore()
+    store.add(all_chunks, all_sources, embeddings)
+    return store
+
+# --- Step 6: RAG query ---
+def query_rag(question, store, client):
+    q_emb = get_embeddings([question], client)
+    results = store.search(q_emb[0], top_k=3)
     
-    return collection
-
-# --- Step 5: RAG query function ---
-def query_rag(question, collection, client):
-    """Retrieves relevant chunks and generates an answer."""
-    q_embedding = get_embeddings([question], client)
-    
-    results = collection.query(query_embeddings=q_embedding, n_results=3)
-    context = "\n\n".join(results["documents"][0])
+    context = "\n\n".join([chunk for chunk, src, score in results])
     
-    prompt = f"""You are an AI Twin that represents a person. Use ONLY the following context to answer the question. 
+    prompt = f"""You are an AI Twin that represents a person. Use ONLY the following context to answer the question.
 If you don't know the answer from the context, say "I don't have that information in my profile."
 
 Context:
@@ -123,31 +111,29 @@ Answer:"""
         )
         return response.strip()
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        return f"Error: {str(e)}"
 
-# --- Global Initialization ---
-print("Initializing HF client...")
+# --- Initialization ---
+print("Starting AI Twin...")
 hf_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN", None)
 hf_client = InferenceClient(token=hf_token)
 
-print("Loading documents...")
 docs, fnames = load_documents()
 print(f"Loaded {len(docs)} documents: {fnames}")
 
-print("Building vector store (using API for embeddings)...")
-kb_collection = build_vector_store(docs, fnames, hf_client)
-print("Vector store ready!")
+vector_store = build_store(docs, fnames, hf_client)
+print("Ready!")
 
-# --- Gradio UI ---
-def load_profile_summary():
+# --- UI ---
+def load_profile():
     try:
         with open(os.path.join(KNOWLEDGE_BASE_DIR, "profile.txt"), "r", encoding="utf-8") as f:
             return f.read()
-    except FileNotFoundError:
+    except:
         return "Profile not found."
 
-def ask_ai_twin(message, chat_history):
-    answer = query_rag(message, kb_collection, hf_client)
+def respond(message, chat_history):
+    answer = query_rag(message, vector_store, hf_client)
     chat_history.append((message, answer))
     return "", chat_history
 
@@ -158,8 +144,7 @@ with gr.Blocks(title="My AI Twin", theme=gr.themes.Soft()) as demo:
     with gr.Row():
         with gr.Column(scale=1):
             gr.Markdown("### 📋 Profile Summary")
-            profile_content = load_profile_summary()
-            gr.Textbox(value=profile_content, label="About Me", interactive=False, lines=15)
+            gr.Textbox(value=load_profile(), label="About Me", interactive=False, lines=15)
         
         with gr.Column(scale=2):
             chatbot = gr.Chatbot(label="Conversation", height=400)
@@ -168,8 +153,8 @@ with gr.Blocks(title="My AI Twin", theme=gr.themes.Soft()) as demo:
                 submit_btn = gr.Button("Submit", variant="primary")
                 clear = gr.Button("Clear")
 
-            msg.submit(ask_ai_twin, [msg, chatbot], [msg, chatbot])
-            submit_btn.click(ask_ai_twin, [msg, chatbot], [msg, chatbot])
+            msg.submit(respond, [msg, chatbot], [msg, chatbot])
+            submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
             clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":

requirements.txt

@@ -1,4 +1,3 @@
-chromadb
 gradio
 huggingface-hub
 numpy