Update model.py

model.py

@@ -1,75 +1,62 @@
 import os
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFaceHub
-from langchain.prompts import PromptTemplate
-from langchain.chains import RetrievalQA
-from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
+from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.document_loaders import TextLoader
+from langchain.text_splitter import CharacterTextSplitter
 from langchain.docstore.document import Document
+from langchain.chains import RetrievalQA
+from langchain_community.llms import HuggingFaceHub
+from langchain.embeddings.base import Embeddings
 
-# Load Hugging Face API token from environment
-HUGGINGFACEHUB_API_TOKEN = os.environ.get("HUGGINGFACEHUB_API_TOKEN")
+# Set safe caching directories to avoid permission denied errors
+os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
+os.environ["HF_HOME"] = "/app/cache"
+os.makedirs("/app/cache", exist_ok=True)
 
-# Embedding model (can be changed to any sentence transformer model)
-embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+# Constants
+DATA_PATH = "/app/data"
+VECTORSTORE_PATH = "/app/vectorstore"
+DOCS_FILENAME = "context.txt"
+EMBEDDING_MODEL_NAME = "sentence-transformers/paraphrase-MiniLM-L6-v2"
 
-# 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."
 
-Context: {context}
-Question: {question}
-Answer:"""
-)
+def load_embedding_model() -> Embeddings:
+    """Initialize and return the HuggingFace embedding model."""
+    return HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
 
-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()
 
-    # Split into smaller chunks
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
-    docs = text_splitter.split_documents(documents)
+def load_documents() -> list[Document]:
+    """Load and split documents into chunks."""
+    loader = TextLoader(os.path.join(DATA_PATH, DOCS_FILENAME))
+    raw_docs = loader.load()
+    splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+    docs = splitter.split_documents(raw_docs)
+    return docs
 
-    # Create FAISS vectorstore
-    vectordb = FAISS.from_documents(docs, embedding_model)
-    vectordb.save_local("vectorstore")
-    return vectordb
 
-def load_vectorstore():
-    """Load existing FAISS vectorstore from disk."""
-    return FAISS.load_local("vectorstore", embedding_model, allow_dangerous_deserialization=True)
+def load_vectorstore() -> FAISS:
+    """Load or create FAISS vectorstore from documents."""
+    vectorstore_file = os.path.join(VECTORSTORE_PATH, "faiss_index")
 
-def get_llm():
-    """Load the HuggingFace Mistral LLM."""
-    return HuggingFaceHub(
-        repo_id="mistralai/Mistral-7B-Instruct-v0.1",
-        model_kwargs={"temperature": 0.5, "max_new_tokens": 512},
-        huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN
-    )
+    embedding_model = load_embedding_model()
+
+    if os.path.exists(vectorstore_file):
+        return FAISS.load_local(vectorstore_file, embedding_model, allow_dangerous_deserialization=True)
+
+    docs = load_documents()
+    vectorstore = FAISS.from_documents(docs, embedding_model)
+    vectorstore.save_local(vectorstore_file)
+    return vectorstore
 
-def build_qa_chain():
-    """Build the full RAG QA chain."""
-    vectordb = load_vectorstore()
-    retriever = vectordb.as_retriever()
-    llm = get_llm()
 
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=retriever,
-        return_source_documents=True,
-        chain_type_kwargs={"prompt": prompt_template}
+def ask_question(query: str) -> str:
+    """Query the vectorstore and return the answer using the language model."""
+    vectorstore = load_vectorstore()
+
+    llm = HuggingFaceHub(
+        repo_id="mistralai/Mistral-7B-Instruct-v0.1",
+        model_kwargs={"temperature": 0.5, "max_tokens": 256},
     )
-    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"]]
-    }
+    qa = RetrievalQA.from_chain_type(llm=llm, retriever=vectorstore.as_retriever())
+    return qa.run(query)