Update app.py

app.py

@@ -99,14 +99,13 @@ class Retriever:
 
     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="valhalla/bart-large-finetuned-squadv1"):
+                 generator_name="facebook/bart-large"):
 
       # generator_name = "valhalla/bart-large-finetuned-squadv1"
       # generator_name = "'vblagoje/bart_lfqa'"
@@ -122,22 +121,24 @@ 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)
+      self.generator_tokenizer = BartTokenizer.from_pretrained(generator_name)
+      self.generator_model = BartForConditionalGeneration.from_pretrained(generator_name).to(device)
       context = self.retriever.retrieve_top_k(question, k=5)
+      # input_text = question + " " + " ".join(context)
 
       input_text = "answer: " + " ".join(context) + " " + question
 
-      inputs = self.generator_tokenizer.encode(input_text, return_tensors='pt', max_length=300, truncation=True).to(device)
-      outputs = self.generator_model.generate(inputs, max_length=300, min_length=2, length_penalty=2.0, num_beams=4, early_stopping=True)
+      inputs = self.generator_tokenizer.encode(input_text, return_tensors='pt', max_length=500, truncation=True).to(device)
+      outputs = self.generator_model.generate(inputs, max_length=150, min_length=2, length_penalty=2.0, num_beams=4, early_stopping=True)
 
       answer = self.generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
       return answer
 
     def extractive_query(self, question):
-      context = self.retriever.retrieve_top_k(question, k=7)
+      context = self.retriever.retrieve_top_k(question, k=15)
+      
       
-      inputs = self.generator_tokenizer(question, ". ".join(context), return_tensors="pt", truncation=True, max_length=150, 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)