Update flask_app.py

flask_app.py

@@ -6,8 +6,10 @@ from langchain.schema import Generation
 import torch
 
 # Load your tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("NLPGenius/KPDastak-llama3-1b")
-model = AutoModelForCausalLM.from_pretrained("NLPGenius/KPDastak-llama3-1b", torch_dtype=torch.bfloat16, device_map="auto",)
+tokenizer = AutoTokenizer.from_pretrained("NLPGenius/GovGPT-llama3")
+model = AutoModelForCausalLM.from_pretrained("NLPGenius/GovGPT-llama3", 
+                                             torch_dtype=torch.bfloat16, 
+                                             device_map="auto",)
 
 
 import transformers
@@ -36,21 +38,53 @@ from langchain.vectorstores import FAISS
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
 embedding_model = HuggingFaceEmbeddings()  # Adjust if you have a specific embedding model
 
-from langchain_core.output_parsers import StrOutputParser
+import os
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.vectorstores import FAISS
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain
 from langchain_core.runnables import RunnablePassthrough
-from langchain_community.document_loaders import Docx2txtLoader
+from langchain_core.output_parsers import StrOutputParser
 
-# Set up the vector store for FAISS
-def create_faiss_index(pdf_path):
-    loader = Docx2txtLoader(pdf_path)
-    data = loader.load()
-    
-    texts = text_splitter.split_documents(data)
-    
+# Function to create FAISS index
+def create_faiss_index(docs):
+    texts = text_splitter.split_documents(docs)
     vector_store = FAISS.from_documents(texts, embedding_model)
+    return vector_store
+
+# Function to fetch and process PDF files
+def fetch_and_process_pdfs(folder_path):
+    pdf_files = [
+        os.path.join(folder_path, file) for file in os.listdir(folder_path) if file.endswith(".pdf")
+    ]
+    
+    all_documents = []
+    for pdf_file in pdf_files:
+        loader = PyPDFLoader(pdf_file)
+        documents = loader.load()
+        all_documents.extend(documents)
+    #print(all_documents)
+    return all_documents
+
+
+# Function to create RAG pipeline
+def create_rag_pipeline(folder_path):
+    # Fetch and process PDF files
+    documents = fetch_and_process_pdfs(folder_path)
+
+    # Create vector store
+    vector_store = create_faiss_index(documents)
     retriever = vector_store.as_retriever()
+
     return retriever
 
+# Usage example
+folder_path = "/kaggle/input/pdf-files/Rules folder - Copy"
+retriever = create_rag_pipeline(folder_path)
+
+
 
 # Define a custom output parser to extract only the answer
 class AnswerOutputParser(BaseOutputParser):
@@ -60,19 +94,21 @@ class AnswerOutputParser(BaseOutputParser):
             return text.split("Answer:")[-1].strip()
         return "I don't know."
 
-# Define the RAG pipeline
-def create_rag_pipeline(pdf_path):
-    retriever = create_faiss_index(pdf_path)
-
-    # Define the prompt template
+# Function to perform model inference
+def model_inference(retriever, question):
+    # Define prompt template
     prompt = PromptTemplate(
         input_variables=["context", "question"],
         template="""\
         You are a helpful assistant. Answer the query accurately by focusing solely on the most relevant parts of the provided context.
+
         1. Identify and use only the sections of the context directly related to the query, ignoring unrelated or extraneous information.
         2. If the query cannot be explicitly answered using the relevant parts of the context, respond with: "The answer is not found in the context provided."
-        
-        Context:
+        3. If you are unsure or do not know the answer, respond with: "I don't know."
+        4. Do not add information, interpretations, or assumptions beyond what is explicitly stated in the context.
+        5. Ensure the response is concise, avoids redundancy, and directly addresses the query.
+
+         Context:
         {context}
         
         Query:
@@ -82,18 +118,20 @@ def create_rag_pipeline(pdf_path):
         """
     )
 
-    # Define the sequence
+
     rag_chain = (
-        RunnableMap({"context": retriever, "question": RunnablePassthrough()})  # Map inputs
-        | prompt  # Use the prompt template
-        | llm  # Apply the LLM
-        | AnswerOutputParser()  # Extract answer from output
-    )
+            {"context": retriever, "question": RunnablePassthrough()}
+            | prompt
+            | llm
+            | AnswerOutputParser() 
+        )
+    response = rag_chain.invoke(question)
+    print("Bot response: ",response)
     return rag_chain
 
-# Usage example
-pdf_path = "Dastak_FullContext.docx"
-rag_pipeline = create_rag_pipeline(pdf_path)
+#response = model_inference(retriever, "Why was a new agriculture policy needed in Khyber Pakhtunkhwa?")
+
+
 
 from flask import Flask, request, jsonify
 from flask_cors import CORS
@@ -106,7 +144,7 @@ CORS(app, supports_credentials=True, allow_headers=["Content-Type"])
 def chatbot_response(query):
     try:
         # Call your actual RAG pipeline here
-        response = rag_pipeline.invoke(query)  # Replace with your actual RAG pipeline call
+        response = model_inference(retriever, query)
         return response
     except Exception as e:
         return f"Error: {str(e)}"