feat: Implement a basic RAG chatbot application using local ChromaDB, HuggingFace embeddings and model, and a Gradio interface.

app.py

@@ -0,0 +1,91 @@
+import gradio as gr
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+import torch
+import os
+from data_cutter import create_db
+
+# Constants
+CHROMA_PATH = "chroma_db"
+MODEL_ID = "Qwen/Qwen2.5-0.5B-Instruct"
+
+print("🚀 Starting app...")
+
+# 1. Initialize/Load Database
+print("🔄 Initializing database from data folder...")
+# We rebuild the DB on startup to ensure it matches the current data
+try:
+    vectorstore = create_db()
+    print("✅ Database created successfully!")
+except Exception as e:
+    print(f"❌ Error creating database: {e}")
+    # Fallback: try to load if exists, though create_db should have handled it
+    if os.path.exists(CHROMA_PATH):
+        print("⚠️ Attempting to load existing database...")
+        embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+        vectorstore = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)
+    else:
+        raise e
+
+# 2. Load AI Model
+print(f"🤖 Loading AI Model ({MODEL_ID})...")
+try:
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+    model = AutoModelForCausalLM.from_pretrained(MODEL_ID)
+    
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=512,
+        device=-1,  # Run on CPU
+        do_sample=True,
+        temperature=0.7,
+        top_p=0.9,
+    )
+    print("✅ AI Model loaded successfully!")
+except Exception as e:
+    print(f"❌ Error loading model: {e}")
+    raise e
+
+def chat_function(message, history):
+    print(f"📨 Received query: {message}")
+    
+    # Search documents
+    results = vectorstore.similarity_search(message, k=3)
+    context = "\n\n".join([doc.page_content for doc in results])
+    
+    # Prepare prompt
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant. Answer the user's question based ONLY on the provided context. If the answer is not in the context, say you don't know."},
+        {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {message}"}
+    ]
+    
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    
+    # Generate response
+    outputs = pipe(prompt)
+    generated_text = outputs[0]['generated_text']
+    
+    # Extract response
+    if "<|im_start|>assistant" in generated_text:
+        response_text = generated_text.split("<|im_start|>assistant")[-1].strip()
+    elif prompt in generated_text:
+        response_text = generated_text.replace(prompt, "").strip()
+    else:
+        response_text = generated_text
+        
+    return response_text
+
+# Create Gradio Interface
+demo = gr.ChatInterface(
+    fn=chat_function,
+    title="RAG Chat with Your Data",
+    description=f"Ask questions about your documents. Powered by {MODEL_ID}.",
+    examples=["What is the main topic?", "Summarize the content."],
+    type="messages"
+)
+
+if __name__ == "__main__":
+    demo.launch()

data_cutter.py

@@ -0,0 +1,63 @@
+from langchain_community.document_loaders import DirectoryLoader, TextLoader
+from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_chroma import Chroma
+from dotenv import load_dotenv
+import os
+
+# Load environment variables from .env file
+load_dotenv()
+
+DATA_PATH = "data"
+CHROMA_PATH = "chroma_db"
+
+def load_documents():
+    loader = DirectoryLoader(DATA_PATH, glob="*.txt", loader_cls=TextLoader)
+    documents = loader.load()
+    return documents
+
+def create_db():
+    documents = load_documents()
+
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size = 1000,
+        chunk_overlap = 500,
+        length_function= len,
+        add_start_index= True,
+    )
+
+    chunks = text_splitter.split_documents(documents)
+
+    print(f"Loaded {len(documents)} document(s)")
+    print(f"Split into {len(chunks)} chunks")
+
+    # Clear existing database if it exists
+    if os.path.exists(CHROMA_PATH):
+        print(f"\nClearing existing database at {CHROMA_PATH}...")
+        import shutil
+        shutil.rmtree(CHROMA_PATH)
+
+    # Create embeddings and vector store
+    print(f"\nCreating ChromaDB vector store with HuggingFace embeddings (all-MiniLM-L6-v2)...")
+    embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
+
+    # Create the vector store from documents
+    vectorstore = Chroma.from_documents(
+        documents=chunks,
+        embedding=embeddings,
+        persist_directory=CHROMA_PATH
+    )
+
+    print(f"✅ Successfully created ChromaDB with {len(chunks)} chunks!")
+    print(f"📁 Database saved to: {CHROMA_PATH}")
+    return vectorstore
+
+if __name__ == "__main__":
+    vectorstore = create_db()
+    
+    # Test the vector store with a simple query
+    print(f"\n🔍 Testing vector store with a sample query...")
+    test_results = vectorstore.similarity_search("Alice", k=3)
+    print(f"Found {len(test_results)} relevant chunks for query 'Alice'")
+    print(f"\nFirst result preview:")
+    print(f"{test_results[0].page_content[:200]}..." if test_results else "No results")

requirements.txt

@@ -0,0 +1,9 @@
+gradio
+python-dotenv
+langchain-huggingface
+langchain-chroma
+langchain-community
+langchain-text-splitters
+transformers
+torch
+