Update app.py

app.py

@@ -67,16 +67,16 @@ class Retriever:
   def load_chunks(self):
     self.text = self.extract_text_from_pdf(self.file_path)
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=300,
+        chunk_size=150,
         chunk_overlap=20,
         length_function=self.token_len,
-        separators=["\n\n", " ", ".", ""]
+        separators=["Section", "\n\n", "\n", ".", " ", ""]
     )
 
     self.chunks = text_splitter.split_text(self.text)
 
   def load_context_embeddings(self):
-    encoded_input = self.context_tokenizer(self.chunks, return_tensors='pt', padding=True, truncation=True, max_length=100).to(device)
+    encoded_input = self.context_tokenizer(self.chunks, return_tensors='pt', padding=True, truncation=True, max_length=300).to(device)
 
     with torch.no_grad():
       model_output = self.context_model(**encoded_input)
@@ -89,20 +89,16 @@ class Retriever:
     encoded_query = self.question_tokenizer(query_prompt, return_tensors="pt", truncation=True, padding=True).to(device)
 
     with torch.no_grad():
-      model_output = self.question_model(**encoded_query)
-      query_vector = model_output.pooler_output
+        model_output = self.question_model(**encoded_query)
+        query_vector = model_output.pooler_output
 
     query_vector_np = query_vector.cpu().numpy()
     D, I = self.index.search(query_vector_np, k)
 
-    retrieved_texts = [self.chunks[i] for i in I[0]]
+    retrieved_texts = [' '.join(self.chunks[i].split('\n')) for i in I[0]]  # Replacing newlines with spaces
 
     scores = [d for d in D[0]]
 
-    # print("Top 5 retrieved texts and their associated scores:")
-    # for idx, (text, score) in enumerate(zip(retrieved_texts, scores)):
-    #     print(f"{idx + 1}. Text: {text} \n   Score: {score:.4f}\n")
-
     return retrieved_texts
 
 class RAG:
@@ -115,22 +111,23 @@ class RAG:
 
       # generator_name = "valhalla/bart-large-finetuned-squadv1"
       # generator_name = "'vblagoje/bart_lfqa'"
-      generator_name = "a-ware/bart-squadv2"
-      
+      # 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)
+
       self.retriever = Retriever(file_path, device, context_model_name, question_model_name)
       self.retriever.load_chunks()
       self.retriever.load_context_embeddings()
 
-    def get_answer(self, question, context):
-      input_text = "context: %s <question for context: %s </s>" % (context,question)
-      features = self.generator_tokenizer([input_text], return_tensors='pt')
-      out = self.generator_model.generate(input_ids=features['input_ids'].to(device), attention_mask=features['attention_mask'].to(device))
-      return self.generator_tokenizer.decode(out[0])
 
-    def query(self, question):
+    def abstractive_query(self, question):
       context = self.retriever.retrieve_top_k(question, k=5)
       # input_text = question + " " + " ".join(context)
 
@@ -144,22 +141,46 @@ class RAG:
       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=15)
+      generator_name = "valhalla/bart-large-finetuned-squadv1"
 
-context_model_name="facebook/dpr-ctx_encoder-single-nq-base"
-context_model_name="facebook/dpr-ctx_encoder-multiset-base"
-question_model_name="facebook/dpr-question_encoder-multiset-base"
+      self.generator_tokenizer = AutoTokenizer.from_pretrained(generator_name)
+      self.generator_model = BartForQuestionAnswering.from_pretrained(generator_name).to(device)
 
-rag = RAG(file_path, device)
+      inputs = self.generator_tokenizer(question, ". ".join(context), return_tensors="pt", truncation=True, max_length=200 , padding="max_length")
+      with torch.no_grad():
+        model_inputs = inputs.to(device)
+        outputs = self.generator_model(**model_inputs)
+
+      answer_start_index = outputs.start_logits.argmax()
+      answer_end_index = outputs.end_logits.argmax()
 
-query = "what is the benefit of using multiple attention heads in mult-head attention?"
+      if answer_end_index < answer_start_index:
+        answer_start_index, answer_end_index = answer_end_index, answer_start_index
 
-print(rag.query(query))
+      print(answer_start_index, answer_end_index)
+
+      predict_answer_tokens = inputs.input_ids[0, answer_start_index : answer_end_index + 1]
+      answer = self.generator_tokenizer.decode(predict_answer_tokens, skip_special_tokens=True)
+      answer = answer.replace('\n', ' ').strip()
+      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)
 
 st.title("RAG Model Query Interface")
 
-query = st.text_area("Enter your question:")
+# offer to ask a question and get an answer. make it so they can ask as many questions as they want
+
+question = st.text_input("Ask a question", "What is another name for self-attention?")
 
-# If a query is given, get the answer
-if query:
-    answer = rag.query(query)
-    st.write(answer)
+if st.button("Ask"):
+  answer = rag.extractive_query(question)
+  st.write(answer)