Updated app.py

app.py

@@ -1,20 +1,23 @@
 import streamlit as st
-# from langchain.document_loaders import PyPDFLoader
 from PyPDF2 import PdfReader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_community.embeddings.spacy_embeddings import SpacyEmbeddings
-from langchain_community.vectorstores import FAISS
-from langchain.chat_models import ChatOllama
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain_community.chat_models import ChatOllama
 from langchain.prompts import ChatPromptTemplate
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
 from langchain.schema.runnable import RunnablePassthrough
 from langchain.schema.output_parser import StrOutputParser
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import os
 
+# Initialize embeddings model
 embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
+# Load the LLaMA model and tokenizer
+model_name = "meta-llama/Llama-2-7b-chat-hf"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
+
 def pdf_read(pdf_doc):
     st.write("Loading PDF")
     text = ""
@@ -38,7 +41,6 @@ def vector_store(text_chunks):
 
 def get_conversational_chain(retriever, user_question):
     with st.spinner("Calling LLM"):
-        llm = ChatOllama(model_name="llama2", temperature=0)
         template = """You are an assistant for question-answering tasks for Retrieval Augmented Generation system for the financial reports such as 10Q and 10K.
         Use the following pieces of retrieved context to answer the question.
         If you don't know the answer, just say that you don't know.
@@ -46,17 +48,22 @@ def get_conversational_chain(retriever, user_question):
         Question: {question}
         Context: {context}
         Answer:
-        """  
+        """
         prompt = ChatPromptTemplate.from_template(template)
+
         # Setup RAG pipeline
         conversation_chain = (
-            {"context": retriever,  "question": RunnablePassthrough()}
+            {"context": retriever, "question": RunnablePassthrough()}
             | prompt
-            | llm
+            | model
             | StrOutputParser()
         )
+        
         st.write("Creating chain")
-        answer = conversation_chain.invoke(user_question)
+        inputs = tokenizer(user_question, return_tensors="pt")
+        outputs = model.generate(**inputs)
+        answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
         st.write(answer)
         st.write("Answer: ", answer)
     return answer