ModularityAI/LLama3Rag

1import yaml2import fitz3import torch4import gradio as gr5from PIL import Image6from langchain.embeddings import HuggingFaceEmbeddings7from langchain.vectorstores import Chroma8from langchain.chains import ConversationalRetrievalChain9from langchain.document_loaders import PyPDFLoader10from langchain.prompts import PromptTemplate11from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline12import spaces13from langchain_text_splitters import CharacterTextSplitter,RecursiveCharacterTextSplitter14 15 16class PDFChatBot:17    def __init__(self, config_path="config.yaml"):18        """19        Initialize the PDFChatBot instance.20 21        Parameters:22            config_path (str): Path to the configuration file (default is "../config.yaml").23        """24        self.processed = False25        self.page = 026        self.chat_history = []27        # Initialize other attributes to None28        self.prompt = None29        self.documents = None30        self.embeddings = None31        self.vectordb = None32        self.tokenizer = None33        self.model = None34        self.pipeline = None35        self.chain = None36        self.chunk_size = 51237        self.overlap_percentage = 5038        self.max_chunks_in_context = 239        self.current_context = None40        self.model_temperatue = 0.541        self.format_seperator="""\n\n--\n\n"""42        self.pipe = None43        #self.chunk_size_slider = chunk_size_slider44 45    def load_embeddings(self):46 47        self.embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")48        print("Embedding model loaded")49 50    def load_vectordb(self):51        overlap = int((self.overlap_percentage/100) * self.chunk_size)52        text_splitter = RecursiveCharacterTextSplitter(53            chunk_size=self.chunk_size,54            chunk_overlap=overlap,55            length_function=len,56            add_start_index=True,57        )58        docs = text_splitter.split_documents(self.documents)59        self.vectordb = Chroma.from_documents(docs, self.embeddings)60        print("Vector store created")61    @spaces.GPU62    def load_tokenizer(self):63        self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3-8B-Instruct")64 65    @spaces.GPU66    def create_organic_pipeline(self):67        self.pipe = pipeline(68            "text-generation",69            model="meta-llama/Meta-Llama-3-8B-Instruct",70            model_kwargs={"torch_dtype": torch.bfloat16},71            device="cuda",72        )73        print("Model pipeline loaded")74 75    def get_organic_context(self, query):76        documents = self.vectordb.similarity_search_with_relevance_scores(query, k=self.max_chunks_in_context)77        context = self.format_seperator.join([doc.page_content for doc, score in documents])78        self.current_context = context79        print("Context Ready")80        print(self.current_context)81    @spaces.GPU82    def create_organic_response(self, history, query):83        self.get_organic_context(query)84        """85        pipe = pipeline(86            "text-generation",87            model="meta-llama/Meta-Llama-3-8B-Instruct",88            model_kwargs={"torch_dtype": torch.bfloat16},89            device="cuda",90        )91        """92        messages = [93            {"role": "system", "content": "From the the contained given below, answer the question of user \n " + self.current_context},94            {"role": "user", "content": query},95        ]96 97        prompt = self.pipe.tokenizer.apply_chat_template(98            messages,99            tokenize=False,100            add_generation_prompt=True101        )102        temp = 0.1103        outputs = self.pipe(104            prompt,105            max_new_tokens=1024,106            do_sample=True,107            temperature=temp,108            top_p=0.9,109        )110        print(outputs)111        return outputs[0]["generated_text"][len(prompt):]112 113 114    def process_file(self, file):115        """116        Process the uploaded PDF file and initialize necessary components: Tokenizer, VectorDB and LLM.117 118        Parameters:119            file (FileStorage): The uploaded PDF file.120        """121        self.documents = PyPDFLoader(file.name).load()122        self.load_embeddings()123        self.load_vectordb()124        self.create_organic_pipeline()125        #self.create_chain()126    @spaces.GPU127    def generate_response(self, history, query, file,chunk_size,chunk_overlap_percentage,model_temperature,max_chunks_in_context):128 129        self.chunk_size = chunk_size130        self.overlap_percentage = chunk_overlap_percentage131        self.model_temperatue = model_temperature132        self.max_chunks_in_context = max_chunks_in_context133 134        if not query:135            raise gr.Error(message='Submit a question')136        if not file:137            raise gr.Error(message='Upload a PDF')138        if not self.processed:139            self.process_file(file)140            self.processed = True141 142 143 144        result = self.create_organic_response(history="",query=query)145        for char in result:146            history[-1][-1] += char147        return history,""148 149    def render_file(self, file,chunk_size,chunk_overlap_percentage,model_temperature,max_chunks_in_context):150        print(chunk_size)151        doc = fitz.open(file.name)152        page = doc[self.page]153        self.chunk_size = chunk_size154        self.overlap_percentage = chunk_overlap_percentage155        self.model_temperatue = model_temperature156        self.max_chunks_in_context = max_chunks_in_context157        pix = page.get_pixmap(matrix=fitz.Matrix(300 / 72, 300 / 72))158        image = Image.frombytes('RGB', [pix.width, pix.height], pix.samples)159        return image160 161    def add_text(self, history, text):162        """163        Add user-entered text to the chat history.164        Parameters:165            history (list): List of chat history tuples.166            text (str): User-entered text.167        Returns:168            list: Updated chat history.169        """170        if not text:171            raise gr.Error('Enter text')172        history.append((text, ''))173        return history