Update main.py

main.py

@@ -1,12 +1,21 @@
-from fastapi import FastAPI, File, UploadFile
+
+
+
+
+
+
+
+from fastapi import FastAPI, File, UploadFile, Form
 import os
 import shutil
 import tempfile
-from langchain.document_loaders import PyPDFLoader
-from langchain_community.llms import CTransformers
-from langchain.chains import LLMChain
-from langchain.prompts import PromptTemplate
+from langchain_community.document_loaders import PyMuPDFLoader
+from LoadLLM import Loadllm
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.chains import ConversationalRetrievalChain
 
+DB_FAISS_PATH = 'vectorstore/db_faiss'
 
 app = FastAPI()
 
@@ -15,7 +24,7 @@ async def home():
     return "API Server Running"
 
 @app.post('/PromptBuddy')
-async def PromptLLM(file: UploadFile = File(...)):
+async def PromptLLM(file: UploadFile = File(...), query: str = Form(...)):
 
     with tempfile.NamedTemporaryFile(delete=False) as temp_file:  # Create temporary file
       temp_file_path = temp_file.name
@@ -23,18 +32,25 @@ async def PromptLLM(file: UploadFile = File(...)):
         shutil.copyfileobj(file.file, f)
     
 
-    loader = PyPDFLoader(temp_file_path)
+    loader = PyMuPDFLoader(file_path= temp_file_path)
     data = loader.load()
 
-    llm = CTransformers(model="llama-2-7b-chat.ggmlv3.q4_1.bin", model_type="llama",
-                            config={'max_new_tokens': 1024,  'context_length': 2048, 'temperature': 0.01})
+    # Create embeddings using Sentence Transformers
+    embeddings = HuggingFaceEmbeddings(model_name='sentence-transformers/all-MiniLM-L6-v2')
+
+    # Create a FAISS vector store and save embeddings
+    db = FAISS.from_documents(data, embeddings)
+    db.save_local(DB_FAISS_PATH)
+
+    # Load the language model
+    llm = Loadllm.load_llm()
+
+    # Create a conversational chain
+    chain = ConversationalRetrievalChain.from_llm(llm=llm, retriever=db.as_retriever())
+
+    result = chain({"question": query, "chat_history": ''})
+    return result['answer']
 
-    template = """Summarise the report {pages}"""
-    prompt_template = PromptTemplate(input_variables=["pages"], template=template)
-    chain = LLMChain(llm=llm, prompt=prompt_template)
 
-    result = chain.run(pages=data[0].page_content)
-    
-    return result