Update app.py

app.py

@@ -1,38 +1,69 @@
-from langchain.document_loaders import TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import OpenAIEmbeddings
-from langchain.vectorstores import FAISS
-from langchain.llms import OpenAI
-from langchain.chains import ConversationalRetrievalChain
-
 import gradio as gr
-import os
+from transformers import AutoTokenizer, AutoModelForCausalLM, AutoModel
+import torch
+import faiss
+import numpy as np
+
+# --- Load TinyLlama Chat Model ---
+llama_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+llama_tokenizer = AutoTokenizer.from_pretrained(llama_id)
+llama_model = AutoModelForCausalLM.from_pretrained(llama_id)
+
+# --- Load Nomic Embedding Model ---
+nomic_id = "nomic-ai/nomic-embed-text-v1"
+nomic_tokenizer = AutoTokenizer.from_pretrained(nomic_id)
+nomic_model = AutoModel.from_pretrained(nomic_id)
+
+# --- FAISS Setup ---
+doc_store = []
+doc_index = None
+
+def embed(texts):
+    inputs = nomic_tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        outputs = nomic_model(**inputs)
+    embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()
+    return embeddings
+
+def add_documents(text_block):
+    global doc_index, doc_store
+    docs = [line.strip() for line in text_block.split("\n") if line.strip()]
+    vectors = embed(docs)
+
+    doc_store.extend(docs)
+
+    if doc_index is None:
+        doc_index = faiss.IndexFlatL2(vectors.shape[1])
+    doc_index.add(vectors)
 
-# Set your API key
-#os.environ["OPENAI_API_KEY"] = "your-openai-key"
+    return f"Added {len(docs)} documents."
 
-# Load and split the document
-loader = TextLoader("your_doc.txt")  # or use PDFLoader for PDFs
-documents = loader.load()
+def chat_with_tinyllama(question):
+    if doc_index is None or len(doc_store) == 0:
+        context = "No documents uploaded yet."
+    else:
+        q_embed = embed([question])
+        D, I = doc_index.search(q_embed, 1)
+        context = doc_store[I[0][0]]
 
-splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-docs = splitter.split_documents(documents)
+    prompt = f"<|system|> You are a helpful assistant.\n<|user|> Context: {context}\nQuestion: {question}\n<|assistant|>"
+    input_ids = llama_tokenizer(prompt, return_tensors="pt").input_ids
+    output = llama_model.generate(input_ids, max_new_tokens=100, do_sample=True)
+    answer = llama_tokenizer.decode(output[0], skip_special_tokens=True)
+    return answer.split("<|assistant|>")[-1].strip()
 
-# Embed and store
-embedding = OpenAIEmbeddings()
-vectorstore = FAISS.from_documents(docs, embedding)
+# --- Gradio Interface ---
+with gr.Blocks() as demo:
+    gr.Markdown("# 🤖 TinyLlama + Nomic Chatbot")
+    with gr.Row():
+        doc_input = gr.Textbox(lines=5, label="Upload Documents (one per line)")
+        upload_btn = gr.Button("Add to Knowledge Base")
 
-# Create retrieval-augmented QA chain
-qa_chain = ConversationalRetrievalChain.from_llm(
-    llm=OpenAI(temperature=0),
-    retriever=vectorstore.as_retriever()
-)
+    with gr.Row():
+        question = gr.Textbox(lines=1, label="Ask a Question")
+        response = gr.Textbox(lines=4, label="TinyLlama Response")
 
-# Memory for chat history
-chat_history = []
+    upload_btn.click(add_documents, inputs=doc_input, outputs=doc_input)
+    question.submit(chat_with_tinyllama, inputs=question, outputs=response)
 
-def chatbot(user_input):
-    global chat_history
-    result = qa_chain({"question": user_input, "chat_history": chat_history})
-    chat_history.append((user_input, result["answer"]))
-    return chat_history
+demo.launch()