Upload 2 files

app.py

@@ -1,69 +1,93 @@
 import os
 import gradio as gr
-from langchain_community.document_loaders import DirectoryLoader, TextLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.vectorstores import Chroma
-from langchain_huggingface import HuggingFaceEndpoint
-from langchain.chains import RetrievalQA
-from langchain_core.prompts import PromptTemplate
+from sentence_transformers import SentenceTransformer
+import chromadb
+from huggingface_hub import InferenceClient
 
 # --- Configuration ---
-# You can set your Hugging Face Token here or as an environment variable
-# os.environ["HUGGINGFACEHUB_API_TOKEN"] = "your_token_here" 
-
 KNOWLEDGE_BASE_DIR = "knowledge_base"
-PERSIST_DIRECTORY = "chroma_db"
+COLLECTION_NAME = "ai_twin_kb"
 
+# --- Step 1: Load documents from knowledge_base/ ---
 def load_documents():
-    """Loads text documents from the knowledge base directory."""
-    loader = DirectoryLoader(KNOWLEDGE_BASE_DIR, glob="*.txt", loader_cls=TextLoader)
-    documents = loader.load()
-    return documents
-
-def create_vector_store(documents):
-    """Splits documents and creates a Chroma vector store."""
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
-    texts = text_splitter.split_documents(documents)
-
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    """Loads all .txt files from the knowledge base directory."""
+    documents = []
+    filenames = []
+    for filename in os.listdir(KNOWLEDGE_BASE_DIR):
+        if filename.endswith(".txt"):
+            filepath = os.path.join(KNOWLEDGE_BASE_DIR, filename)
+            with open(filepath, "r", encoding="utf-8", errors="ignore") as f:
+                content = f.read().strip()
+                if content:
+                    documents.append(content)
+                    filenames.append(filename)
+    return documents, filenames
+
+# --- 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):
+        end = start + chunk_size
+        chunks.append(text[start:end])
+        start += chunk_size - overlap
+    return chunks
+
+# --- Step 3: Build vector store ---
+def build_vector_store(documents, filenames):
+    """Creates embeddings and stores them in ChromaDB."""
+    model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+    
+    client = chromadb.Client()
+    # Delete existing collection if it exists
+    try:
+        client.delete_collection(COLLECTION_NAME)
+    except:
+        pass
+    collection = client.create_collection(name=COLLECTION_NAME)
     
-    # Check if vector store already exists to avoid re-creating it every time? 
-    # For this assignment, re-creating it ensures latest data is used.
-    if os.path.exists(PERSIST_DIRECTORY):
-        try:
-            # simple cleanup for fresh clear start (optional for production but good for dev)
-            import shutil
-            shutil.rmtree(PERSIST_DIRECTORY)
-        except:
-            pass
-            
-    vector_store = Chroma.from_documents(texts, embeddings, persist_directory=PERSIST_DIRECTORY)
-    return vector_store
-
-def setup_rag_chain(vector_store):
-    """Sets up the RAG chain with a retrieval capability."""
-    # Using a free endpoint model. 
-    # 'mistralai/Mistral-7B-Instruct-v0.2' is a good choice, but requires a token.
-    # 'google/flan-t5-large' is another option.
-    # We'll use a generic reliable one or let the user input their token/model in a real scenario.
-    # For the assignment, let's try to use a model that might work with the free tier or a locally downloadable one if needed.
-    # However, running local LLM is heavy.
-    # Let's assume the user has a token or we use a very small model.
-    # If no token is found, this might fail or warn.
+    all_chunks = []
+    all_ids = []
+    all_metadata = []
+    chunk_id = 0
     
-    llm = HuggingFaceEndpoint(
-        repo_id="mistralai/Mistral-7B-Instruct-v0.2", 
-        task="text-generation",
-        max_new_tokens=512,
-        do_sample=False,
-        repetition_penalty=1.03,
+    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
+    
+    # Generate embeddings
+    embeddings = model.encode(all_chunks).tolist()
+    
+    # Add to ChromaDB
+    collection.add(
+        documents=all_chunks,
+        embeddings=embeddings,
+        ids=all_ids,
+        metadatas=all_metadata
     )
     
-    retriever = vector_store.as_retriever(search_kwargs={"k": 3})
+    return collection, model
+
+# --- Step 4: RAG query function ---
+def query_rag(question, collection, embed_model, llm_client):
+    """Retrieves relevant chunks and generates an answer."""
+    # Embed the question
+    q_embedding = embed_model.encode([question]).tolist()
+    
+    # Retrieve top 3 relevant chunks
+    results = collection.query(query_embeddings=q_embedding, n_results=3)
+    
+    # Build context from retrieved documents
+    context = "\n\n".join(results["documents"][0])
     
-    prompt_template = """Use the following pieces of context to answer the question at the end. 
-If you don't know the answer, just say that you don't know, don't try to make up an answer.
+    # Create prompt
+    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:
 {context}
@@ -71,74 +95,65 @@ Context:
 Question: {question}
 
 Answer:"""
-    PROMPT = PromptTemplate(
-        template=prompt_template, input_variables=["context", "question"]
-    )
-
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=retriever,
-        return_source_documents=True,
-        chain_type_kwargs={"prompt": PROMPT}
-    )
-    return qa_chain
+    
+    # Generate response using Hugging Face Inference API
+    try:
+        response = llm_client.text_generation(
+            prompt,
+            max_new_tokens=512,
+            temperature=0.3,
+            repetition_penalty=1.1
+        )
+        return response.strip()
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
 
 # --- Global Initialization ---
 print("Loading documents...")
-docs = load_documents()
-print(f"Loaded {len(docs)} documents.")
-
-print("Creating vector store...")
-vector_db = create_vector_store(docs)
-print("Vector store created.")
-
-print("Setting up RAG chain...")
-try:
-    rag_chain = setup_rag_chain(vector_db)
-    print("RAG chain setup complete.")
-except Exception as e:
-    print(f"Error setting up RAG chain (likely missing HF Token): {e}")
-    rag_chain = None
-
-def ask_ai_twin(question):
-    if not rag_chain:
-        return "Error: RAG Chain not initialized. Please check your Hugging Face Token."
-    
-    result = rag_chain.invoke({"query": question})
-    return result["result"]
+docs, fnames = load_documents()
+print(f"Loaded {len(docs)} documents: {fnames}")
+
+print("Building vector store...")
+kb_collection, embedding_model = build_vector_store(docs, fnames)
+print("Vector store ready.")
+
+print("Initializing LLM client...")
+hf_token = os.environ.get("HUGGINGFACEHUB_API_TOKEN", None)
+llm = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.2", token=hf_token)
+print("LLM client ready.")
 
 # --- Gradio UI ---
 def load_profile_summary():
     try:
-        with open(os.path.join(KNOWLEDGE_BASE_DIR, "profile.txt"), "r") as f:
+        with open(os.path.join(KNOWLEDGE_BASE_DIR, "profile.txt"), "r", encoding="utf-8") as f:
             return f.read()
     except FileNotFoundError:
         return "Profile not found."
 
-with gr.Blocks(title="My AI Twin") as demo:
-    gr.Markdown("# My AI Twin")
+def ask_ai_twin(message, chat_history):
+    answer = query_rag(message, kb_collection, embedding_model, llm)
+    chat_history.append((message, answer))
+    return "", chat_history
+
+with gr.Blocks(title="My AI Twin", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🤖 My AI Twin")
     gr.Markdown("Ask me anything about my professional background, skills, and projects!")
     
     with gr.Row():
         with gr.Column(scale=1):
-            gr.Markdown("### Profile Summary")
+            gr.Markdown("### 📋 Profile Summary")
             profile_content = load_profile_summary()
-            gr.Textbox(value=profile_content, label="About Me", interactive=False, lines=10)
+            gr.Textbox(value=profile_content, label="About Me", interactive=False, lines=15)
         
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="Conversation")
-            msg = gr.Textbox(label="Ask a question")
-            submit_btn = gr.Button("Submit")
-            clear = gr.Button("Clear")
-
-            def respond(message, chat_history):
-                bot_message = ask_ai_twin(message)
-                chat_history.append((message, bot_message))
-                return "", chat_history
-
-            msg.submit(respond, [msg, chatbot], [msg, chatbot])
-            submit_btn.click(respond, [msg, chatbot], [msg, chatbot])
+            chatbot = gr.Chatbot(label="Conversation", height=400)
+            msg = gr.Textbox(label="Ask a question", placeholder="e.g. What are my skills?")
+            with gr.Row():
+                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])
             clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":

requirements.txt

@@ -1,8 +1,3 @@
-langchain
-langchain-community
-langchain-huggingface
-langchain-text-splitters
-langchain-core
 chromadb
 sentence-transformers
 gradio