Update app.py

app.py

@@ -1,194 +1,68 @@
+# app.py
 import os
-import faiss
-import numpy as np
-import gradio as gr
-
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from PyPDF2 import PdfReader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from langchain.chains import RetrievalQA
+from langchain.llms import HuggingFaceHub
+import gradio as gr
 
-# -----------------------------
-# CONFIG
-# -----------------------------
-DATA_PATH = "Docs"
-TOP_K = 3
-LLM_MODEL = "google/flan-t5-base"
-
-# -----------------------------
-# LOAD MODELS
-# -----------------------------
-embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL)
-llm_model = AutoModelForSeq2SeqLM.from_pretrained(LLM_MODEL)
-
-# -----------------------------
-# FILE LOADER
-# -----------------------------
-def read_file(path):
-    if path.endswith(".txt") or path.endswith(".md"):
-        with open(path, "r", encoding="utf-8") as f:
-            return f.read()
-    elif path.endswith(".pdf"):
-        reader = PdfReader(path)
-        text = ""
-        for page in reader.pages:
-            text += page.extract_text() or ""
-        return text
-    return ""
+# --------- Configuration ---------
+DOCS_FOLDER = "Docs"  # Folder containing PDF files
+HF_API_KEY = os.getenv("HF_TOKEN")  # HuggingFace API token
 
-def load_docs(folder):
+# --------- Load and process PDFs ---------
+def load_pdfs(folder):
     texts = []
-    if not os.path.exists(folder):
-        return texts
-
     for file in os.listdir(folder):
-        path = os.path.join(folder, file)
-        try:
-            txt = read_file(path)
-            if txt.strip():
-                texts.append(txt)
-        except:
-            continue
+        if file.endswith(".pdf"):
+            pdf_path = os.path.join(folder, file)
+            reader = PdfReader(pdf_path)
+            for page in reader.pages:
+                texts.append(page.extract_text())
     return texts
 
-# -----------------------------
-# CHUNKING
-# -----------------------------
-def chunk_text(text, size=300, overlap=50):
-    words = text.split()
-    return [" ".join(words[i:i + size]) for i in range(0, len(words), size - overlap)]
-
-# -----------------------------
-# BUILD VECTOR DB
-# -----------------------------
-def build_index(docs):
-    chunks = []
-    for doc in docs:
-        chunks.extend(chunk_text(doc))
-
-    if not chunks:
-        return None, []
-
-    embeddings = embedding_model.encode(chunks)
-    dim = embeddings.shape[1]
-
-    index = faiss.IndexFlatL2(dim)
-    index.add(np.array(embeddings))
-
-    return index, chunks
-
-# -----------------------------
-# RETRIEVE
-# -----------------------------
-def retrieve(query, index, chunks, k=TOP_K):
-    q_embed = embedding_model.encode([query])
-    D, I = index.search(np.array(q_embed), k)
-    return [chunks[i] for i in I[0]]
-
-# -----------------------------
-# GENERATE ANSWER (WITH MEMORY)
-# -----------------------------
-def generate_answer(query, contexts, history):
-    context = "\n\n".join(contexts)
-
-    history_text = ""
-    for h in history[-6:]:
-        history_text += f"{h['role']}: {h['content']}\n"
-
-    prompt = f"""
-You are an expert AI/ML assistant.
-
-Conversation:
-{history_text}
-
-Context:
-{context}
-
-Question:
-{query}
-
-Answer clearly:
-"""
+raw_texts = load_pdfs(DOCS_FOLDER)
+
+# Split into smaller chunks for embeddings
+text_splitter = RecursiveCharacterTextSplitter(
+    chunk_size=1000,
+    chunk_overlap=50
+)
+docs = text_splitter.split_text(" ".join(raw_texts))
+
+# --------- Create embeddings and vectorstore ---------
+embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+vectorstore = FAISS.from_texts(docs, embedding=embeddings)
+
+# --------- Setup LLM & Retrieval QA chain ---------
+llm = HuggingFaceHub(
+    repo_id="google/flan-t5-large", 
+    model_kwargs={"temperature": 0, "max_length": 512}, 
+    huggingfacehub_api_token=HF_API_KEY
+)
+qa_chain = RetrievalQA.from_chain_type(
+    llm=llm, 
+    chain_type="stuff", 
+    retriever=vectorstore.as_retriever()
+)
+
+# --------- Gradio interface ---------
+def answer_query(query):
+    return qa_chain.run(query)
 
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-    outputs = llm_model.generate(**inputs, max_new_tokens=200)
-
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# -----------------------------
-# TOOL RECOMMENDER (AGENT MODE)
-# -----------------------------
-def tool_recommender(query):
-    prompt = f"""
-You are an AI architect.
-
-Suggest best AI/ML tools for:
-
-{query}
-
-Give:
-- Tools
-- Why
-- Use cases
-"""
-
-    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
-    outputs = llm_model.generate(**inputs, max_new_tokens=150)
-
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# -----------------------------
-# INIT
-# -----------------------------
-docs = load_docs(DATA_PATH)
-index, chunks = build_index(docs)
-
-# -----------------------------
-# MAIN CHAT PIPELINE
-# -----------------------------
-def rag_chat(query, history):
-    history = history or []
-
-    if index is None:
-        history.append({"role": "user", "content": query})
-        history.append({"role": "assistant", "content": "No documents found"})
-        return history
-
-    retrieved = retrieve(query, index, chunks)
-
-    # Agent decision
-    if "recommend" in query.lower() or "best tool" in query.lower():
-        answer = tool_recommender(query)
-    else:
-        answer = generate_answer(query, retrieved, history)
-
-    history.append({"role": "user", "content": query})
-    history.append({"role": "assistant", "content": answer})
-
-    return history
-
-# -----------------------------
-# UI HANDLER
-# -----------------------------
-def respond(message, chat_history):
-    chat_history = chat_history or []
-    updated_history = rag_chat(message, chat_history)
-    return "", updated_history
-
-# -----------------------------
-# UI (CHATGPT STYLE)
-# -----------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("## AI/ML Conversational RAG + Tool Recommender")
-
-    chatbot = gr.Chatbot(type="messages")
-    msg = gr.Textbox(placeholder="Ask about AI tools, ML, companies...")
-
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-
-# -----------------------------
-# RUN
-# -----------------------------
-if __name__ == "__main__":
-    demo.launch()
+    gr.Markdown("## PDF Document RAG QA System")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Enter your question:")
+    submit = gr.Button("Ask")
+    
+    def chat_fn(user_input, chat_history):
+        answer = answer_query(user_input)
+        chat_history.append((user_input, answer))
+        return chat_history, ""
+    
+    submit.click(chat_fn, inputs=[msg, chatbot], outputs=[chatbot, msg])
+
+demo.launch(server_name="0.0.0.0", server_port=7860)