import streamlit as st import torch import numpy as np import faiss import PyPDF2 import os import langchain from transformers import DPRContextEncoder, DPRContextEncoderTokenizer, DPRQuestionEncoder, DPRQuestionEncoderTokenizer, BartForQuestionAnswering from transformers import BartForConditionalGeneration, BartTokenizer, AutoTokenizer from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain.document_loaders import PyPDFLoader from streamlit import runtime runtime.exists() 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=150, chunk_overlap=20, length_function=self.token_len, 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=300).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 = [' '.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 = "'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) self.retriever = Retriever(file_path, device, context_model_name, question_model_name) self.retriever.load_chunks() self.retriever.load_context_embeddings() def abstractive_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 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") 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() if answer_end_index < answer_start_index: answer_start_index, answer_end_index = answer_end_index, answer_start_index 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") # 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 st.button("Ask"): answer = rag.extractive_query(question) st.write(answer)