taskswithcode/semantic_search

1from scipy.spatial.distance import cosine2import argparse3import json4import os5import openai6import pdb7 8def read_text(input_file):9    arr = open(input_file).read().split("\n")10    return arr[:-1]11 12 13class OpenAIQnAModel:14    def __init__(self):15        self.debug = False16        self.q_model_name = None17        self.d_model_name = None18        self.skip_key = True19        print("In OpenAI API constructor")20 21 22    def init_model(self,model_name = None):23        #print("OpenAI: Init model",model_name)24        openai.api_key = os.getenv("OPENAI_API_KEY")25        if (openai.api_key == None):26            openai.api_key = ""27            print("API key not set")28 29        if (len(openai.api_key) == 0 and not self.skip_key):30                print("Open API key not set")31        32        if (model_name is None):33            self.d_model_name = "text-search-ada-doc-001"34        else:35            self.d_model_name = model_name36        self.q_model_name = self.construct_query_model_name(self.d_model_name)37        print(f"OpenAI: Init model complete :query model {self.q_model_name} doc:{self.d_model_name}")38 39    def construct_query_model_name(self,d_model_name):40        return d_model_name.replace('-doc-','-query-')41 42 43    def compute_embeddings(self,input_file_name,input_data,is_file):44        if (len(openai.api_key) == 0 and not self.skip_key):45                print("Open API key not set")46                return [],[]47        #print("In compute embeddings after key check")48        in_file = input_file_name.split('/')[-1]49        in_file = self.d_model_name + '_' +  '.'.join(in_file.split('.')[:-1]) + "_search.json"50        cached = False51        try:52            fp = open(in_file)53            cached = True54            embeddings = json.load(fp)55            q_embeddings = [embeddings[0]]56            d_embeddings = embeddings[1:]57            print("Using cached embeddings")58        except:59            pass60            61        texts = read_text(input_data) if is_file == True else input_data62        queries = [texts[0]]63        docs = texts[1:]64 65        if (not cached):66            print(f"Computing embeddings for {input_file_name} and query model {self.q_model_name}")67            query_embeds = openai.Embedding.create(68                input=queries,69                model=self.q_model_name70            )71            print(f"Computing embeddings for {input_file_name} and doc model {self.q_model_name}")72            doc_embeds = openai.Embedding.create(73                input=docs,74                model=self.d_model_name75            )76            q_embeddings = []77            d_embeddings = []78            for i in range(len(query_embeds['data'])):79                q_embeddings.append(query_embeds['data'][i]['embedding'])80            for i in range(len(doc_embeds['data'])):81                d_embeddings.append(doc_embeds['data'][i]['embedding'])82        if (not cached):83            embeddings = q_embeddings + d_embeddings84            with open(in_file,"w") as fp:85                json.dump(embeddings,fp)86        return texts,(q_embeddings,d_embeddings)87 88    def output_results(self,output_file,texts,embeddings,main_index = 0):89        # Calculate cosine similarities90        # Cosine similarities are in [-1, 1]. Higher means more similar91        query_embeddings = embeddings[0]92        doc_embeddings = embeddings[1]93        cosine_dict = {}94        queries = [texts[0]]95        docs = texts[1:]96        if (self.debug):97            print("Total sentences",len(texts))98        for i in range(len(docs)):99            cosine_dict[docs[i]] = 1 - cosine(query_embeddings[0], doc_embeddings[i])100 101        if (self.debug):102            print("Input sentence:",texts[main_index])103        sorted_dict = dict(sorted(cosine_dict.items(), key=lambda item: item[1],reverse = True))104        if (self.debug):105            for key in sorted_dict:106                print("Cosine similarity with  \"%s\" is: %.3f" % (key, sorted_dict[key]))107        if (output_file is not None):108            with open(output_file,"w") as fp:109                fp.write(json.dumps(sorted_dict,indent=0))110        return sorted_dict111 112 113 114if __name__ == '__main__':115        parser = argparse.ArgumentParser(description='OpenAI model for document search embeddings ',formatter_class=argparse.ArgumentDefaultsHelpFormatter)116        parser.add_argument('-input', action="store", dest="input",required=True,help="Input file with sentences")117        parser.add_argument('-output', action="store", dest="output",default="output.txt",help="Output file with results")118        parser.add_argument('-model', action="store", dest="model",default="text-search-ada-doc-001",help="model name")119 120        results = parser.parse_args()121        obj = OpenAIQnAModel()122        obj.init_model(results.model)123        texts, embeddings = obj.compute_embeddings(results.input,results.input,is_file = True)124        results = obj.output_results(results.output,texts,embeddings)125