Update functions.py

functions.py

@@ -94,10 +94,12 @@ initial_qa_template = (
     "answer the question: {question}\n.\n"
 )
 
-
+###################### Functions #######################################################################################
 
 @st.experimental_singleton(suppress_st_warning=True)
 def load_models():
+
+    '''Load and cache all the models to be used'''
     q_model = ORTModelForSequenceClassification.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
     ner_model = AutoModelForTokenClassification.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
     kg_model = AutoModelForSeq2SeqLM.from_pretrained("Babelscape/rebel-large")
@@ -128,7 +130,9 @@ def load_asr_model(asr_model_name):
 #     return sbert
 
 @st.experimental_singleton(suppress_st_warning=True)
-def process_corpus(corpus, tok, chunk_size=200, overlap=50):
+def process_corpus(corpus, tok, title, embeddings, chunk_size=200, overlap=50):
+
+    '''Process text for Semantic Search'''
     
     pinecone.init(api_key="2d1e8029-2d84-4724-9f7c-a4f0f5ae908a", environment="us-west1-gcp")
 
@@ -137,10 +141,19 @@ def process_corpus(corpus, tok, chunk_size=200, overlap=50):
 
     texts = text_splitter.split_text(corpus)
 
-    return texts
+    docsearch = Pinecone.from_texts(
+        texts,
+        embeddings,
+        index_name = index_id,
+        namespace = f'{title}-earnings',
+        metadatas = [
+        {'source':i} for i in range(len(texts))]
+    )
+
+    return docsearch
     
 @st.experimental_memo(suppress_st_warning=True)
-def embed_text(query,corpus,title,embedding_model,chain_type='stuff'):
+def embed_text(query,corpus,title,embedding_model,emb_tok,chain_type='stuff'):
     
     '''Embed text and generate semantic search scores'''
 
@@ -156,15 +169,9 @@ def embed_text(query,corpus,title,embedding_model,chain_type='stuff'):
         
         embeddings = HuggingFaceEmbeddings(model_name=f'sentence-transformers/{embedding_model}')
         
-    docsearch = Pinecone.from_texts(
-        corpus,
-        embeddings,
-        index_name = index_id,
-        namespace = f'{title}-earnings',
-        metadatas = [
-        {'source':i} for i in range(len(texts))]
-    )
-
+    title = title[0]
+    docsearch = process_corpus(corpus,embed_tok,title, embeddings)
+        
     docs = docsearch.similarity_search_with_score(query, k=3, namespace = f'{title}-earnings')
 
     docs = [d[0] for d in docs]
@@ -186,8 +193,14 @@ def embed_text(query,corpus,title,embedding_model,chain_type='stuff'):
 
     elif chain_type == 'refine':
 
-    
-    return hits
+        initial_qa_prompt = PromptTemplate(
+    input_variables=["context_str", "question"], template=initial_qa_template
+)
+        chain = load_qa_chain(OpenAI(temperature=0), chain_type="refine", return_refine_steps=False,
+                     question_prompt=initial_qa_prompt, refine_prompt=refine_prompt)
+        answer = chain({"input_documents": docs, "question": query}, return_only_outputs=True)
+
+        return answer['output_text']
     
 # @st.experimental_memo(suppress_st_warning=True)
 # def embed_text(query,corpus,embedding_model):
@@ -304,7 +317,7 @@ def clean_text(text):
        
 @st.experimental_memo(suppress_st_warning=True)
 def chunk_long_text(text,threshold,window_size=3,stride=2):
-    '''Preprocess text and chunk for semantic search and sentiment analysis'''
+    '''Preprocess text and chunk for sentiment analysis'''
     
     #Convert cleaned text into sentences
     sentences = sent_tokenize(text)