MachineLearningReply/search_mlReply

1import streamlit as st2 3from utils.config import document_store_configs, model_configs4from haystack import Pipeline5from haystack.schema import Answer6from haystack.document_stores import BaseDocumentStore7from haystack.document_stores import InMemoryDocumentStore, OpenSearchDocumentStore, WeaviateDocumentStore8from haystack.nodes import EmbeddingRetriever, FARMReader, PromptNode, PreProcessor9#from haystack.nodes import TextConverter, FileTypeClassifier, PDFToTextConverter10from milvus_haystack import MilvusDocumentStore11#Use this file to set up your Haystack pipeline and querying12 13@st.cache_resource(show_spinner=False)14def start_preprocessor_node():15    print('initializing preprocessor node')16    processor = PreProcessor(17        clean_empty_lines= True,18        clean_whitespace=True,19        clean_header_footer=True,20        #remove_substrings=None,21        split_by="word",22        split_length=100,23        split_respect_sentence_boundary=True,24        #split_overlap=0,25        #max_chars_check= 10_00026    )27    return processor28    #return docs29 30@st.cache_resource(show_spinner=False)31def start_document_store(type: str):32    #This function starts the documents store of your choice based on your command line preference33    print('initializing document store')34    if type == 'inmemory':35        document_store = InMemoryDocumentStore(use_bm25=True, embedding_dim=384)36        '''37        documents = [38            {39                'content': "Pi is a super dog",40                'meta': {'name': "pi.txt"}41            },42            {43                'content': "The revenue of siemens is 5 milion Euro",44                'meta': {'name': "siemens.txt"}45            },46        ]47        document_store.write_documents(documents)48        '''49    elif type == 'opensearch':50        document_store = OpenSearchDocumentStore(scheme = document_store_configs['OPENSEARCH_SCHEME'], 51                                                 username = document_store_configs['OPENSEARCH_USERNAME'], 52                                                 password = document_store_configs['OPENSEARCH_PASSWORD'],53                                                 host = document_store_configs['OPENSEARCH_HOST'],54                                                 port = document_store_configs['OPENSEARCH_PORT'],55                                                 index = document_store_configs['OPENSEARCH_INDEX'],56                                                 embedding_dim = document_store_configs['OPENSEARCH_EMBEDDING_DIM'])57    elif type == 'weaviate':58        document_store = WeaviateDocumentStore(host = document_store_configs['WEAVIATE_HOST'],59                                                port = document_store_configs['WEAVIATE_PORT'],60                                                index = document_store_configs['WEAVIATE_INDEX'],61                                                embedding_dim = document_store_configs['WEAVIATE_EMBEDDING_DIM'])62    elif type == 'milvus':63        document_store = MilvusDocumentStore(uri = document_store_configs['MILVUS_URI'],64                                            index = document_store_configs['MILVUS_INDEX'],65                                            embedding_dim = document_store_configs['MILVUS_EMBEDDING_DIM'],66                                            return_embedding=True)67    return document_store68 69# cached to make index and models load only at start70@st.cache_resource(show_spinner=False)71def start_retriever(_document_store: BaseDocumentStore):72    print('initializing retriever')73    retriever = EmbeddingRetriever(document_store=_document_store,74                                   embedding_model=model_configs['EMBEDDING_MODEL'],75                                   top_k=5)76    #77 78    #_document_store.update_embeddings(retriever)79    return retriever80 81 82@st.cache_resource(show_spinner=False)83def start_reader():84    print('initializing reader')85    reader = FARMReader(model_name_or_path=model_configs['EXTRACTIVE_MODEL'])86    return reader87 88 89 90# cached to make index and models load only at start91@st.cache_resource(show_spinner=False)92def start_haystack_extractive(_document_store: BaseDocumentStore, _retriever: EmbeddingRetriever, _reader: FARMReader):93    print('initializing pipeline')94    pipe = Pipeline()95    pipe.add_node(component=_retriever, name="Retriever", inputs=["Query"])96    pipe.add_node(component= _reader, name="Reader", inputs=["Retriever"])97    return pipe98 99@st.cache_resource(show_spinner=False)100def start_haystack_rag(_document_store: BaseDocumentStore, _retriever: EmbeddingRetriever, openai_key):101    prompt_node = PromptNode(default_prompt_template="deepset/question-answering", 102                             model_name_or_path=model_configs['GENERATIVE_MODEL'],103                             api_key=openai_key,104                             max_length=500)105    pipe = Pipeline()106 107    pipe.add_node(component=_retriever, name="Retriever", inputs=["Query"])108    pipe.add_node(component=prompt_node, name="PromptNode", inputs=["Retriever"])109 110    return pipe111 112#@st.cache_data(show_spinner=True)113def query(_pipeline, question):114    params = {}115    results = _pipeline.run(question, params=params)116    return results117 118def initialize_pipeline(task, document_store, retriever, reader, openai_key = ""):119    if task == 'extractive':120        return start_haystack_extractive(document_store, retriever, reader)121    elif task == 'rag':122        return start_haystack_rag(document_store, retriever, openai_key)123 124 125