samwoof/KalbeCDT-GenomicsAnswerer

1from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor2import os3import re4from typing import Dict, List, Tuple5import warnings6 7from langchain_chroma import Chroma8from langchain_huggingface.llms import HuggingFacePipeline9from langchain_huggingface import ChatHuggingFace, HuggingFaceEmbeddings, HuggingFaceEndpoint10from langchain_core.messages.base import BaseMessage11from langchain_core.messages import HumanMessage, SystemMessage, AIMessage12from langchain.docstore.document import Document13 14from prompts import MAIN_SYSTEM_PROMPT15 16CITATIONS_REGEX = r"(\b\d{2}\_\d{2}\b)"17 18 19# TODO: DOCUMENT AND ADD TYPE HINTS TO ALL FUNCTIONS & CLASSES20class Store:21    def __init__(22            self, 23            name: str,24            embedding_model: str="jinaai/jina-embeddings-v2-base-en",25            presist_dir: str="./chroma_langchain_db",26            doc_k=427        ):28        29        self.embedding_func = HuggingFaceEmbeddings(model_name=embedding_model,model_kwargs={"trust_remote_code":True})30        self.name = name31        self.persist_dir = presist_dir32        self.store = None33        self.doc_k = doc_k34 35    def setup(self):36        if not os.path.isdir(self.persist_dir): 37            warnings.warn(f"Vector store directory {self.persist_dir} does not exist, Creating...") 38 39        self.store = Chroma(40            collection_name=self.name,41            embedding_function=self.embedding_func,42            persist_directory=self.persist_dir43        )44    45    def _get_doc_ids(self, docs: List[Document]) -> List[str]:46        doc_ids = []47        for doc in docs:48            doc_ids.append(f"{os.path.basename(doc.metadata['source'])}_{doc.metadata['page']}")49        return doc_ids50    51    def add_docs(self, docs: List[Document]):52        doc_ids = self._get_doc_ids(docs)53        # self.store.add_documents(documents=docs, ids=doc_ids)54        with ThreadPoolExecutor(max_workers=5) as exe:55            exe.submit(self.store.add_documents, documents=docs, ids=doc_ids)56    57    def delete_docs(self, ids: List[str]):58        self.store.delete(ids=ids)59    60    def similarity_search(self, query: str):61        return self.store.similarity_search(query, k=self.doc_k)62 63 64class Answerer:65    def __init__(66            self,67            vec_store: Store,68            model="NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",69            use_api=True,70            temperature=0.05,71            top_p=0.7,72            max_tokens=2048,73        ):74            self.store = vec_store75 76            if not isinstance(model, str):77                self.model = model78                return79 80            if use_api:81                llm = HuggingFaceEndpoint(82                    repo_id=model,83                    model_kwargs={"max_length":max_tokens},84                    max_new_tokens=max_tokens,85                    temperature=temperature,86                    top_p=top_p,87                    huggingfacehub_api_token=os.environ["HUGGINGFACEHUB_API_TOKEN"],88                )89            else:90                llm = HuggingFacePipeline.from_model_id(91                model_id=model,92                task="text-generation",93                pipeline_kwargs={94                    "max_new_tokens": max_tokens,95                    "temperature":temperature,96                    "top_p": top_p97                    },98                )99            100            self.model = ChatHuggingFace(llm=llm)101 102    @staticmethod103    def update_history(query, history):104        history.append({"role":"user", "content": query})105 106        history_langchain = []107        for msg in history:108            if msg['role'] == "user":109                history_langchain.append(HumanMessage(content=msg['content']))110            elif msg['role'] == "assistant":111                history_langchain.append(AIMessage(content=msg['content']))112            elif msg['role'] == "system":113                history_langchain.append(SystemMessage(content=msg['content']))114        115        return history_langchain, history116 117    # TODO: Perhaps make it so it does a search everytime it gets a query? is that better? leaving for future me to handle.118    def answer_with_search(self, query: str, ctx_docs: List[Document]=None, show_cits: bool=True) -> Tuple[List[Dict], List[Document], List[str]]:119        # TODO: Include the tables extracted120        121        search_results = ctx_docs122        if ctx_docs is None:123            search_results = self.store.similarity_search(query)124        125        citation_mapping = self.store.store.get()126 127        # NOTE: 😭😭😭😭128        #search_results_str = "\n".join([129        #    f"=== ID: 'CTX_{citation_mapping[os.path.basename(res.metadata['source'])+str(res.metadata['page'])]}' START ===\n{res.page_content}\n=== ID: 'CTX_{citation_mapping[os.path.basename(res.metadata['source'])+str(res.metadata['page'])]}' END ===" for res in search_results])130        #file_names = set([os.path.basename(res.metadata['source']) for res in search_results])131        search_results_str = "\n\n".join([res.page_content for res in search_results])132 133        system_prompt = MAIN_SYSTEM_PROMPT.format(context=search_results_str)134        history = [135            SystemMessage(content=system_prompt),136            HumanMessage(content=query)137        ]138        result = self.model.invoke(history)139        citations = [res.group() for res in re.finditer(CITATIONS_REGEX, result.content, re.MULTILINE)]140        cits_pages = set([int(c.split("_")[0])-1 for c in citations])141        citations_pages_ids = []142 143        cits = ""144        for c in cits_pages:145            try:146                cits += f"{c+1:0>2}_xx *{citation_mapping['ids'][c]}*\n"147                citations_pages_ids.append(citation_mapping['ids'][c])148            except IndexError:149                cits += f"{c+1} - N/A\n"150 151        history = [152            {"role":"system", "content": system_prompt},153            {"role":"user", "content": query},154            {"role":"assistant", "content": result.content + (("\n\n**Pages Cited:**\n" + cits) if show_cits else "")}155        ]156 157        return history, search_results, citations_pages_ids158 159    def answer_without_search(self, query: str, history: List[Dict]):160        history_langchain, history = self.update_history(query, history)161        result = self.model.invoke(history_langchain)162        history.append({"role":"assistant", "content": result.content})163 164        return history