Update model.py

model.py

@@ -1,63 +1,75 @@
 import os
-from PyPDF2 import PdfReader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
-from langchain_community.docstore.in_memory import InMemoryDocstore
+from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFaceHub
-from langchain.chains import RetrievalQA
 from langchain.prompts import PromptTemplate
-import uuid
-import faiss
+from langchain.chains import RetrievalQA
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+from langchain.docstore.document import Document
+
+# Load Hugging Face API token from environment
+HUGGINGFACEHUB_API_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
 
-vectorstore = None
+# Embedding model (can be changed to any sentence transformer model)
+embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-def load_vectorstore(pdf_path):
-    global vectorstore
+# Prompt template for Mistral
+prompt_template = PromptTemplate(
+    input_variables=["context", "question"],
+    template="""You are an intelligent assistant. Use the context below to answer the question.
+If the answer is not contained in the context, say "I don't know."
 
-    reader = PdfReader(pdf_path)
-    text = "".join([page.extract_text() or "" for page in reader.pages])
-    splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
-    chunks = splitter.split_text(text)
+Context: {context}
+Question: {question}
+Answer:"""
+)
 
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
-    dim = len(embeddings.embed_query("test"))
-    index = faiss.IndexFlatL2(dim)
+def create_vectorstore(doc_path: str = "data/docs.txt"):
+    """Create or load FAISS vectorstore from the given document."""
+    loader = TextLoader(doc_path)
+    documents = loader.load()
 
-    vectorstore = FAISS(
-        embedding_function=embeddings,
-        index=index,
-        docstore=InMemoryDocstore({}),
-        index_to_docstore_id={}
-    )
-    uuids = [str(uuid.uuid4()) for _ in chunks]
-    vectorstore.add_texts(chunks, ids=uuids)
+    # Split into smaller chunks
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
+    docs = text_splitter.split_documents(documents)
 
+    # Create FAISS vectorstore
+    vectordb = FAISS.from_documents(docs, embedding_model)
+    vectordb.save_local("vectorstore")
+    return vectordb
 
-def ask_question(query):
-    global vectorstore
-    if not vectorstore:
-        return "Please upload and index a document first."
+def load_vectorstore():
+    """Load existing FAISS vectorstore from disk."""
+    return FAISS.load_local("vectorstore", embedding_model, allow_dangerous_deserialization=True)
 
-    llm = HuggingFaceHub(
+def get_llm():
+    """Load the HuggingFace Mistral LLM."""
+    return HuggingFaceHub(
         repo_id="mistralai/Mistral-7B-Instruct-v0.1",
-        huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
-        model_kwargs={"temperature": 0.7, "max_length": 512}
+        model_kwargs={"temperature": 0.5, "max_new_tokens": 512},
+        huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
     )
 
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-    prompt = PromptTemplate(
-        template="Use the context to answer the question:
-Context: {context}
-Question: {question}
-Answer:",
-        input_variables=["context", "question"]
-    )
+def build_qa_chain():
+    """Build the full RAG QA chain."""
+    vectordb = load_vectorstore()
+    retriever = vectordb.as_retriever()
+    llm = get_llm()
 
-    chain = RetrievalQA.from_chain_type(
+    qa_chain = RetrievalQA.from_chain_type(
         llm=llm,
         retriever=retriever,
-        return_source_documents=False,
-        chain_type_kwargs={"prompt": prompt}
+        return_source_documents=True,
+        chain_type_kwargs={"prompt": prompt_template}
     )
-    return chain({"query": query})["result"]
+    return qa_chain
+
+def ask_question(query: str) -> dict:
+    """Handle a single user query."""
+    chain = build_qa_chain()
+    result = chain({"query": query})
+    return {
+        "answer": result["result"],
+        "sources": [doc.metadata.get("source", "unknown") for doc in result["source_documents"]]
+    }