initial commit

app.py

@@ -0,0 +1,163 @@
+import streamlit as st
+import torch
+import numpy as np
+import faiss
+import PyPDF2
+import os
+
+from transformers import BertTokenizer, BertModel
+from transformers import DPRContextEncoder, DPRContextEncoderTokenizer, DPRQuestionEncoder, DPRQuestionEncoderTokenizer, BartForQuestionAnswering
+from transformers import BartForConditionalGeneration, BartTokenizer, AutoTokenizer
+
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain import text_splitter
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import PyPDFLoader
+
+device = torch.device("cpu")
+if torch.cuda.is_available():
+   print("Training on GPU")
+   device = torch.device("cuda:0")
+
+file_url = "https://arxiv.org/pdf/1706.03762.pdf"
+file_path = "assets/attention.pdf"
+
+if not os.path.exists('assets'):
+    os.mkdir('assets')
+
+if not os.path.isfile(file_path):
+    os.system(f'curl -o {file_path} {file_url}')
+else:
+    print("File already exists!")
+
+class Retriever:
+
+  def __init__(self, file_path, device, context_model_name, question_model_name):
+    self.file_path = file_path
+    self.device = device
+
+    self.context_tokenizer = DPRContextEncoderTokenizer.from_pretrained(context_model_name)
+    self.context_model = DPRContextEncoder.from_pretrained(context_model_name).to(device)
+
+    self.question_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(question_model_name)
+    self.question_model = DPRQuestionEncoder.from_pretrained(question_model_name).to(device)
+
+  def token_len(self, text):
+    tokens = self.context_tokenizer.encode(text)
+    return len(tokens)
+
+  def extract_text_from_pdf(self, file_path):
+    with open(file_path, 'rb') as file:
+        reader = PyPDF2.PdfReader(file)
+        text = ''
+        for page in reader.pages:
+            text += page.extract_text()
+    return text
+
+  def get_text(self):
+    with open(self.file_path, 'rb') as file:
+        reader = PyPDF2.PdfReader(file)
+        text = ''
+        for page in reader.pages:
+            text += page.extract_text()
+    return text
+
+  def load_chunks(self):
+    self.text = self.extract_text_from_pdf(self.file_path)
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=300,
+        chunk_overlap=20,
+        length_function=self.token_len,
+        separators=["\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)
+
+    with torch.no_grad():
+      model_output = self.context_model(**encoded_input)
+      self.token_embeddings = model_output.pooler_output.cpu().detach().numpy()
+
+    self.index = faiss.IndexFlatL2(self.token_embeddings.shape[1])
+    self.index.add(self.token_embeddings)
+
+  def retrieve_top_k(self, query_prompt, k=10):
+    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
+
+    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]]
+
+    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:
+    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 = "'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)
+
+      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):
+      context = self.retriever.retrieve_top_k(question, k=5)
+      # input_text = question + " " + " ".join(context)
+
+      input_text = "answer: " + " ".join(context) + " " + question
+
+      print(input_text)
+
+      inputs = self.generator_tokenizer.encode(input_text, return_tensors='pt', max_length=1024, 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
+
+
+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"
+
+rag = RAG(file_path, device)
+
+st.title("RAG Model Query Interface")
+
+query = st.text_input("Enter your question:")
+
+# If a query is given, get the answer
+if query:
+    answer = rag.query(query)
+    st.write(f"Answer: {answer}")
+
+if __name__ == "__main__":
+    # This is used when running locally. Can be removed if deploying to a server.
+    st.run()

requirements.txt

@@ -0,0 +1,8 @@
+streamlit
+torch
+numpy
+faiss-cpu
+PyPDF2
+transformers
+langchain
+