Update app.py

app.py

@@ -23,6 +23,10 @@ PROMPT = """/
 
 """
 
+tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2")
+model = AutoModel.from_pretrained("sentence-transformers/all-mpnet-base-v2")
+reranker_model = CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-6")
+
 def load_documents(file_paths):
     documents = []
     for path in file_paths:
@@ -31,11 +35,11 @@ def load_documents(file_paths):
     return documents
 
 class DocumentIndexer:
-    def __init__(self, documents):
+    def __init__(self, documents, tokenizer, model):
         self.documents = documents
         self.bm25 = BM25Okapi([doc.split() for doc in documents])
-        self.tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-mpnet-base-v2")
-        self.model = AutoModel.from_pretrained("sentence-transformers/all-mpnet-base-v2")
+        self.tokenizer = tokenizer
+        self.model = model
         self.index = self.create_faiss_index()
 
     def create_faiss_index(self):
@@ -63,9 +67,8 @@ class DocumentIndexer:
         return [self.documents[i] for i in indices[0]]
 
 class Reranker:
-    def __init__(self, model_name="cross-encoder/ms-marco-TinyBERT-L-6"):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model = CrossEncoder(model_name)
+    def __init__(self, reranker):
+        self.model = reranker
 
     def rank(self, query, documents):
         pairs = [(query, doc) for doc in documents]
@@ -105,8 +108,8 @@ def chatbot_interface(query, history):
     file_paths = ["company.txt", "Base.txt"]  
     documents = load_documents(file_paths)
     
-    indexer = DocumentIndexer(documents)
-    reranker = Reranker()
+    indexer = DocumentIndexer(documents, tokenizer, model)
+    reranker = Reranker(reranker_model)
 
     chatbot = QAChatbot(indexer, reranker)
     answer = chatbot.generate_answer(query)