import shutil import streamlit as st st.set_page_config( page_title="RAG Configuration", page_icon="🤖", layout="wide", initial_sidebar_state="collapsed" ) import re import os import spire.pdf import fitz from src.Databases import * from langchain.text_splitter import * from sentence_transformers import SentenceTransformer, CrossEncoder from langchain_community.llms import HuggingFaceHub from langchain_huggingface import HuggingFaceEmbeddings from transformers import (AutoFeatureExtractor, AutoModel, AutoImageProcessor) from llama_index.embeddings.huggingface import HuggingFaceEmbedding class SentenceTransformerEmbeddings: """ Wrapper Class for SentenceTransformer Class """ def __init__(self, model_name: str): """ Initiliases a Sentence Transformer """ self.model = SentenceTransformer(model_name) def embed_documents(self, texts): """ Returns a list of embeddings for the given texts. """ return self.model.encode(texts, convert_to_tensor=True).tolist() def embed_query(self, text): """ Returns a list of embeddings for the given text. """ return self.model.encode(text, convert_to_tensor=True).tolist() @st.cache_resource(show_spinner=False) def settings(): return HuggingFaceEmbedding(model_name="BAAI/bge-base-en") @st.cache_resource(show_spinner=False) def pine_embedding_model(): return SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2") # 784 dimension + euclidean @st.cache_resource(show_spinner=False) def weaviate_embedding_model(): return SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") @st.cache_resource(show_spinner=False) def load_image_model(model): extractor = AutoFeatureExtractor.from_pretrained(model) im_model = AutoModel.from_pretrained(model) return extractor, im_model @st.cache_resource(show_spinner=False) def load_bi_encoder(): return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L12-v2", model_kwargs={"device": "cpu"}) @st.cache_resource(show_spinner=False) def pine_embedding_model(): return SentenceTransformerEmbeddings(model_name="all-mpnet-base-v2") # 784 dimension + euclidean @st.cache_resource(show_spinner=False) def weaviate_embedding_model(): return SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2") @st.cache_resource(show_spinner=False) def load_cross(): return CrossEncoder("cross-encoder/ms-marco-TinyBERT-L-2-v2", max_length=512, device="cpu") @st.cache_resource(show_spinner=False) def pine_cross_encoder(): return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-12-v2", max_length=512, device="cpu") @st.cache_resource(show_spinner=False) def weaviate_cross_encoder(): return CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2", max_length=512, device="cpu") @st.cache_resource(show_spinner=False) def load_chat_model(): template = ''' You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question accurately. If the question is not related to the context, just answer 'I don't know'. Question: {question} Context: {context} Answer: ''' return HuggingFaceHub( repo_id="mistralai/Mistral-7B-Instruct-v0.1", model_kwargs={"temperature": 0.5, "max_length": 64, "max_new_tokens": 512, "query_wrapper_prompt": template} ) @st.cache_resource(show_spinner=False) def load_q_model(): return HuggingFaceHub( repo_id="mistralai/Mistral-7B-Instruct-v0.3", model_kwargs={"temperature": 0.5, "max_length": 64, "max_new_tokens": 512} ) @st.cache_resource(show_spinner=False) def load_image_model(model): extractor = AutoFeatureExtractor.from_pretrained(model) im_model = AutoModel.from_pretrained(model) return extractor, im_model @st.cache_resource(show_spinner=False) def load_nomic_model(): return AutoImageProcessor.from_pretrained("nomic-ai/nomic-embed-vision-v1.5"), AutoModel.from_pretrained("nomic-ai/nomic-embed-vision-v1.5", trust_remote_code=True) @st.cache_resource(show_spinner=False) def vector_database_prep(file): def data_prep(file): def findWholeWord(w): return re.compile(r'\b{0}\b'.format(re.escape(w)), flags=re.IGNORECASE).search file_name = file.name pdf_file_path = os.path.join(os.getcwd(), 'pdfs', file_name) image_folder = os.path.join(os.getcwd(), f'figures_{file_name}') if not os.path.exists(image_folder): os.makedirs(image_folder) # everything down here is wrt pages dir print('1. folder made') with spire.pdf.PdfDocument() as doc: doc.LoadFromFile(pdf_file_path) images = [] for page_num in range(doc.Pages.Count): page = doc.Pages[page_num] for image_num in range(len(page.ImagesInfo)): imageFileName = os.path.join(image_folder, f'figure-{page_num}-{image_num}.png') image = page.ImagesInfo[image_num] image.Image.Save(imageFileName) images.append({ "image_file_name": imageFileName, "image": image }) print('2. image extraction done') image_info = [] for image_file in os.listdir(image_folder): if image_file.endswith('.png'): image_info.append({ "image_file_name": image_file[:-4], "image": Image.open(os.path.join(image_folder, image_file)), "pg_no": int(image_file.split('-')[1]) }) print('3. temporary') figures = [] with fitz.open(pdf_file_path) as pdf_file: data = "" for page in pdf_file: text = page.get_text() if not (findWholeWord('table of contents')(text) or findWholeWord('index')(text)): data += text data = data.replace('}', '-') data = data.replace('{', '-') print('4. Data extraction done') hs = [] for i in image_info: src = i['image_file_name'] + '.png' headers = {'_': []} header = '_' page = pdf_file[i['pg_no']] texts = page.get_text('dict') for block in texts['blocks']: if block['type'] == 0: for line in block['lines']: for span in line['spans']: if 'bol' in span['font'].lower() and not span['text'].isnumeric(): header = span['text'] print("header: ", header) headers[header] = [header] else: headers[header].append(span['text']) try: if findWholeWord('fig')(span['text']): i['image_file_name'] = span['text'] figures.append(span['text'].split('fig')[-1]) elif findWholeWord('figure')(span['text']): i['image_file_name'] = span['text'] figures.append(span['text'].lower().split('figure')[-1]) else: pass except re.error: pass if not i['image_file_name'].endswith('.png'): s = i['image_file_name'] + '.png' i['image_file_name'] = s os.rename(os.path.join(image_folder, src), os.path.join(image_folder, i['image_file_name'])) hs.append({"image": i, "header": headers}) print('5. header and figures done') figure_contexts = {} for fig in figures: figure_contexts[fig] = [] for page_num in range(len(pdf_file)): page = pdf_file[page_num] texts = page.get_text('dict') for block in texts['blocks']: if block['type'] == 0: for line in block['lines']: for span in line['spans']: if findWholeWord(fig)(span['text']): print('figure mention: ', span['text']) figure_contexts[fig].append(span['text']) print('6. Figure context collected') contexts = [] for h in hs: context = "" for q in h['header'].values(): context += "".join(q) s = pytesseract.image_to_string(h['image']['image']) qwea = context + '\n' + s if len(s) != 0 else context contexts.append(( h['image']['image_file_name'], qwea, h['image']['image'] )) print('7. Overall context collected') image_content = [] for fig in figure_contexts: for c in contexts: if findWholeWord(fig)(c[0]): s = c[1] + '\n' + "\n".join(figure_contexts[fig]) s = str("\n".join( [ "".join([h for h in i.strip() if h.isprintable()]) for i in s.split('\n') if len(i.strip()) != 0 ] )) image_content.append(( c[0], s, c[2] )) print('8. Figure context added') return data, image_content # Vector Database objects extractor, i_model = st.session_state['extractor'], st.session_state['image_model'] pinecone_embed = st.session_state['pinecone_embed'] weaviate_embed = st.session_state['weaviate_embed'] vb1 = UnifiedDatabase('vb1', 'lancedb/rag') vb1.model_prep(extractor, i_model, weaviate_embed, RecursiveCharacterTextSplitter(chunk_size=1330, chunk_overlap=35)) vb2 = UnifiedDatabase('vb2', 'lancedb/rag') vb2.model_prep(extractor, i_model, pinecone_embed, RecursiveCharacterTextSplitter(chunk_size=1330, chunk_overlap=35)) vb_list = [vb1, vb2] data, image_content = data_prep(file) for vb in vb_list: vb.upsert(data) vb.upsert(image_content) # image_cont = dict[image_file_path, context, PIL] return vb_list os.environ["HUGGINGFACEHUB_API_TOKEN"] = st.secrets["HUGGINGFACEHUB_API_TOKEN"] os.environ["LANGCHAIN_PROJECT"] = st.secrets["LANGCHAIN_PROJECT"] os.environ["OPENAI_API_KEY"] = st.secrets["GPT_KEY"] st.session_state['pdf_file'] = [] st.session_state['vb_list'] = [] st.session_state['Settings.embed_model'] = settings() st.session_state['processor'], st.session_state['vision_model'] = load_nomic_model() st.session_state['bi_encoder'] = load_bi_encoder() st.session_state['chat_model'] = load_chat_model() st.session_state['cross_model'] = load_cross() st.session_state['q_model'] = load_q_model() st.session_state['extractor'], st.session_state['image_model'] = load_image_model("google/vit-base-patch16-224-in21k") st.session_state['pinecone_embed'] = pine_embedding_model() st.session_state['weaviate_embed'] = weaviate_embedding_model() st.title('Multi-modal RAG based LLM for Information Retrieval') st.subheader('Converse with our Chatbot') st.markdown('Enter a pdf file as a source.') uploaded_file = st.file_uploader("Choose an pdf document...", type=["pdf"], accept_multiple_files=False) if uploaded_file is not None: with open(uploaded_file.name, mode='wb') as w: w.write(uploaded_file.getvalue()) if not os.path.exists(os.path.join(os.getcwd(), 'pdfs')): os.makedirs(os.path.join(os.getcwd(), 'pdfs')) shutil.move(uploaded_file.name, os.path.join(os.getcwd(), 'pdfs')) st.session_state['pdf_file'] = uploaded_file.name with st.spinner('Extracting'): vb_list = vector_database_prep(uploaded_file) st.session_state['vb_list'] = vb_list st.switch_page('pages/rag.py')