pipeline refactor with custom pineconer retriever

app_pinecode.py

@@ -4,6 +4,14 @@ import time
 from config import Config
 from pinecode_quieries import PinecodeProposalQueries
 
+import logging
+
+
+logging.basicConfig(format="%(levelname)s - %(name)s -  %(message)s", level=logging.WARNING)
+logging.getLogger("haystack").setLevel(logging.INFO)
+
+logging.info("This is a test log ..")
+
 extractive_query = PinecodeProposalQueries (es_host = Config.es_host, es_index = Config.proposals_index,
                                   es_user = Config.es_user, es_password = Config.es_password,
                                   reader_name_or_path = Config.reader_model_name_or_path,

pinecode_quieries.py

@@ -1,10 +1,34 @@
+
 from abc import ABC, abstractmethod
 from haystack.nodes import BM25Retriever, FARMReader
 from haystack.document_stores import ElasticsearchDocumentStore
-from haystack.pipelines import ExtractiveQAPipeline
+from haystack.pipelines import ExtractiveQAPipeline, DocumentSearchPipeline
 from haystack.document_stores import PineconeDocumentStore
 from haystack.nodes import EmbeddingRetriever
 
+import json
+import logging
+import os
+import shutil
+import sys
+import uuid
+from json import JSONDecodeError
+from pathlib import Path
+from typing import List, Optional
+
+import pandas as pd
+import pinecone
+import streamlit as st
+
+from haystack import BaseComponent, Document
+from haystack.document_stores import PineconeDocumentStore
+from haystack.nodes import (
+    EmbeddingRetriever,
+    FARMReader
+)
+from haystack.pipelines import ExtractiveQAPipeline, Pipeline
+from sentence_transformers import SentenceTransformer
+
 import certifi
 import datetime
 import requests
@@ -12,6 +36,36 @@ from base64 import b64encode
 
 ca_certs=certifi.where()        
 
+class PineconeRetriever(BaseComponent):
+    outgoing_edges = 1
+
+    def __init__(self, sentence_transformer_name: str, api_key:str, environment: str, index_name: str): 
+        # a small subset of the component's parameters is sent in an event after applying filters defined in haystack.telemetry.NonPrivateParameters
+        self.sts_model = SentenceTransformer(sentence_transformer_name)
+        pinecone.init(api_key = api_key, environment=environment)
+        self.index = pinecone.Index(index_name)
+
+    def run(self, query: str, top_k: Optional[int]):
+        # process the inputs
+        vector_embeddings = self.sts_model.encode(query).tolist()
+        response = self.index.query([vector_embeddings], top_k=top_k, include_metadata=True)
+        docs = [
+            Document(
+                content=d["metadata"]['content'],
+                meta={'title': d["metadata"]['title'], 
+                      'page': d["metadata"]['page'],
+                      'source': d["metadata"]['source']
+                      }
+                )
+                for d in response["matches"]
+        ]
+        output = {"documents": docs, "query": query}
+        return output, "output_1"
+
+    def run_batch(self, queries: List[str], top_k: Optional[int]):
+        return {}, "output_1"
+
+
 class DocumentQueries(ABC):
     
     @abstractmethod
@@ -28,34 +82,47 @@ class PinecodeProposalQueries(DocumentQueries):
     def _initialize_pipeline(self, es_host, es_index, es_user, es_password, reader = None):
         if reader is not None:
             self.reader = reader
-        self.es_host = es_host
-        self.es_user = es_user
-        self.es_password = es_password
+        
+        #pinecone.init(api_key=es_password, environment="us-east1-gcp")
+        index_name = "semantic-text-search"
+        
         self.document_store = PineconeDocumentStore(
             api_key=es_password,
             environment = "us-east1-gcp",
-            index="semantic-text-search",
+            index=index_name,
             similarity="cosine",
             embedding_dim=384
-        )        
-        #self.retriever = BM25Retriever(document_store = self.document_store)
-        self.retriever = EmbeddingRetriever(
-            document_store=self.document_store,
-            #embedding_model="multi-qa-distilbert-dot-v1",
-            embedding_model = "sentence-transformers/msmarco-MiniLM-L6-cos-v5",
-            model_format="sentence_transformers"
         )
+        self.pipe = Pipeline()
+        pinecone_retriever = PineconeRetriever("sentence-transformers/multi-qa-MiniLM-L6-cos-v1", 
+                                               es_password, "us-east1-gcp", 
+                                               index_name)
+        self.pipe.add_node(component=pinecone_retriever, name="Retriever", inputs=["Query"])
+        self.pipe.add_node(component=self.reader, name="Reader", inputs=["Retriever"])
+               
+        # #self.retriever = BM25Retriever(document_store = self.document_store)
+        # self.retriever = EmbeddingRetriever(
+        #     document_store=self.document_store,
+        #     #embedding_model="multi-qa-distilbert-dot-v1",
+        #     embedding_model = "sentence-transformers/msmarco-MiniLM-L6-cos-v5",
+        #     model_format="sentence_transformers"
+        # )
+        
+        # retriever_model = "sentence-transformers/multi-qa-mpnet-base-dot-v1"
+
+
+        
         #self.document_store.update_embeddings(self.retriever, update_existing_embeddings=False)
-        self.pipe = ExtractiveQAPipeline(self.reader, self.retriever)
+        #self.pipe = ExtractiveQAPipeline (reader = self.reader, retriever = self.retriever)
+        #self.pipe = DocumentSearchPipeline(self.retriever)
 
     def search_by_query(self, query : str, retriever_top_k: int, reader_top_k: int, es_index: str = None) :
         #self.document_store.update_embeddings(self.retriever, update_existing_embeddings=False)
 
-        #self.log.write_log(query, "hfspace-informecomision")
         #if es_index is not None:
             #self._initialize_pipeline(self.es_host, es_index, self.es_user, self.es_password)
-        #params = {"Retriever": {"top_k": retriever_top_k}, "Reader": {"top_k": reader_top_k}}
-        params = {"Retriever": {"top_k": retriever_top_k}}
+        params = {"Retriever": {"top_k": retriever_top_k}, "Reader": {"top_k": reader_top_k}}
+        #params = {"Retriever": {"top_k": retriever_top_k}}
         prediction = self.pipe.run( query = query, params = params)      
         return prediction["answers"]
 

requirements.txt

@@ -2,4 +2,3 @@ transformers==4.19.2
 torch==1.10.2
 #farm-haystack==1.5.0
 farm-haystack[pinecone]==1.5.0
-pinecone-client<=2.0.10