Update app.py

app.py

@@ -97,30 +97,24 @@ class Retriever:
 
     retrieved_texts = [' '.join(self.chunks[i].split('\n')) for i in I[0]]  # Replacing newlines with spaces
 
-    scores = [d for d in D[0]]
-
     return retrieved_texts
 
+
 class RAG:
     def __init__(self,
                  file_path,
                  device,
                  context_model_name="facebook/dpr-ctx_encoder-multiset-base",
                  question_model_name="facebook/dpr-question_encoder-multiset-base",
-                 generator_name="facebook/bart-large"):
+                 generator_name="valhalla/bart-large-finetuned-squadv1"):
 
       # generator_name = "valhalla/bart-large-finetuned-squadv1"
       # generator_name = "'vblagoje/bart_lfqa'"
       # generator_name = "a-ware/bart-squadv2"
 
-      self.generator_tokenizer = BartTokenizer.from_pretrained(generator_name)
-      self.generator_model = BartForConditionalGeneration.from_pretrained(generator_name).to(device)
-
-      # generator_name = "MaRiOrOsSi/t5-base-finetuned-question-answering"
-      # generator_name = "t5-small"
-
-      # self.generator_tokenizer = T5Tokenizer.from_pretrained(generator_name)
-      # self.generator_model = T5ForConditionalGeneration.from_pretrained(generator_name)
+      generator_name = "valhalla/bart-large-finetuned-squadv1"
+      self.generator_tokenizer = AutoTokenizer.from_pretrained(generator_name)
+      self.generator_model = BartForQuestionAnswering.from_pretrained(generator_name).to(device)
 
       self.retriever = Retriever(file_path, device, context_model_name, question_model_name)
       self.retriever.load_chunks()
@@ -128,8 +122,9 @@ class RAG:
 
 
     def abstractive_query(self, question):
+      self.generator_tokenizer = BartTokenizer.from_pretrained(self.generator_name)
+      self.generator_model = BartForConditionalGeneration.from_pretrained(self.generator_name).to(device)
       context = self.retriever.retrieve_top_k(question, k=5)
-      # input_text = question + " " + " ".join(context)
 
       input_text = "answer: " + " ".join(context) + " " + question
 
@@ -141,12 +136,9 @@ class RAG:
 
     def extractive_query(self, question):
       context = self.retriever.retrieve_top_k(question, k=15)
-      generator_name = "valhalla/bart-large-finetuned-squadv1"
-
-      self.generator_tokenizer = AutoTokenizer.from_pretrained(generator_name)
-      self.generator_model = BartForQuestionAnswering.from_pretrained(generator_name).to(device)
-
-      inputs = self.generator_tokenizer(question, ". ".join(context), return_tensors="pt", truncation=True, max_length=200 , padding="max_length")
+      
+      
+      inputs = self.generator_tokenizer(question, ". ".join(context), return_tensors="pt", truncation=True, max_length=300 , padding="max_length")
       with torch.no_grad():
         model_inputs = inputs.to(device)
         outputs = self.generator_model(**model_inputs)
@@ -163,11 +155,9 @@ class RAG:
       answer = answer.replace('$', '')
 
       return answer
-    
+
 context_model_name="facebook/dpr-ctx_encoder-single-nq-base"
 question_model_name = "facebook/dpr-question_encoder-single-nq-base"
-# context_model_name="facebook/dpr-ctx_encoder-multiset-base"
-# question_model_name="facebook/dpr-question_encoder-multiset-base"
 
 rag = RAG(file_path, device)