Update app.py

app.py

@@ -1,51 +1,84 @@
 from datasets import load_dataset
-from sentence_transformers import SentenceTransformer
-import faiss
 import numpy as np
 import gradio as gr
 import chromadb
 
-# Load a small subset (10,000 rows)
-dataset = load_dataset("wiki40b", "en", split="train[:1000]")
+from transformers import AutoModel, AutoTokenizer, pipeline
+import torch
+import chromadb
 
-# Extract only text
-docs = [d["text"] for d in dataset]
+# Initialize ChromaDB client
+chroma_client = chromadb.PersistentClient(path="./chroma_db")  # Stores data persistently
+collection = chroma_client.get_or_create_collection(name="wikipedia_docs")
 
-print("Loaded dataset with", len(docs), "documents.")
+# Load the BAAI embedding model
+model_name = "BAAI/bge-base-en"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModel.from_pretrained(model_name)
 
-# Load embedding model
-embed_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+def get_embedding(text):
+    """Generate embeddings using BAAI/bge-base-en."""
+    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    return outputs.last_hidden_state[:, 0, :].numpy().tolist()  # Take CLS token embedding
 
-# Convert texts to embeddings
-embeddings = embed_model.encode(docs, show_progress_bar=True)
+# Load LLaMA Model (Meta LLaMA 2)
+llama_pipe = pipeline("text-generation", model="meta-llama/Llama-2-7b-chat-hf")
 
+# Load a small subset (10,000 rows)
+#dataset = load_dataset("wiki40b", "en", split="train[:1000]")
 
-# Initialize ChromaDB client
-chroma_client = chromadb.PersistentClient(path="./chroma_db")  # Stores data persistently
-collection = chroma_client.get_or_create_collection(name="wikipedia_docs")
+# Extract only text
+#docs = [d["text"] for d in dataset]
+docs = ["Machine learning is a field of AI...", "Neural networks are inspired by the brain..."]
 
-# Store embeddings in ChromaDB
-for i, (doc, embedding) in enumerate(zip(docs, embeddings)):
-    collection.add(
-        ids=[str(i)],  # Unique ID for each doc
-        embeddings=[embedding.tolist()],  # Convert numpy array to list
-        documents=[doc]
-    )
+#print("Loaded dataset with", len(docs), "documents.")
+
+# ✅ Step 2: Embed and Store in ChromaDB
+for i, doc in enumerate(docs):
+    embedding = get_embedding(doc)
+    collection.add(ids=[str(i)], embeddings=[embedding], documents=[doc])
 
 print("Stored embeddings in ChromaDB!")
 
+# Store embeddings in ChromaDB
+#for i, (doc, embedding) in enumerate(zip(docs, embeddings)):
+#    collection.add(
+#        ids=[str(i)],  # Unique ID for each doc
+#        embeddings=[embedding.tolist()],  # Convert numpy array to list
+#        documents=[doc]
+#    )
+
+
 # Search function using ChromaDB
-def search_wikipedia(query, top_k=3):
-    query_embedding = embed_model.encode([query]).tolist()
-    results = collection.query(
-        query_embeddings=query_embedding, 
-        n_results=top_k
-    )
-    return "\n\n".join(results["documents"][0])  # Return top results
+#def search_wikipedia(query, top_k=3):
+#    query_embedding = embed_model.encode([query]).tolist()
+#    results = collection.query(
+#        query_embeddings=query_embedding, 
+#        n_results=top_k
+    
+    #return "\n\n".join(results["documents"][0])  # Return top results
+#    return results["documents"][0]  # Return top results
+
+# Function to search ChromaDB & generate response
+def query_llama(user_input):
+    query_embedding = get_embedding(user_input)
+    results = collection.query(query_embeddings=[query_embedding], n_results=3)
+
+    if not results["documents"]:
+        return "No relevant documents found."
+
+    context = " ".join(results["documents"][0])
+    prompt = f"Using this context, answer the question: {user_input}\nContext: {context}"
+
+    response = llama_pipe(prompt, max_length=200)
+    return f"**LLaMA Response:** {response[0]['generated_text']}\n\n**Retrieved Docs:** {context}"
+
 
 # Gradio Interface
 iface = gr.Interface(
-    fn=search_wikipedia, 
+    fn=query_llama, 
     inputs="text", 
     outputs="text", 
     title="Wikipedia Search RAG",