Rakesh30/Sentence_Embedding-App

1import gradio as gr2import pickle3import os4from datasets import load_dataset5from gradio.components import Label6from InstructorEmbedding import INSTRUCTOR7import heapq8from sklearn.metrics.pairwise import cosine_similarity9import nltk10from nltk.corpus import stopwords11from nltk.tokenize import word_tokenize, sent_tokenize12from nltk.stem import WordNetLemmatizer13import pandas as pd14 15dataset = load_dataset("SandipPalit/Movie_Dataset")16 17model = INSTRUCTOR('hkunlp/instructor-xl')18 19def getSimilarity(sentences_a,sentences_b):20  embeddings_a = pickle.load(open(os.getcwd()+"/temp.pkl",'rb'))21  embeddings_b = model.encode(sentences_b)22  similarities = cosine_similarity(embeddings_a,embeddings_b)23  return similarities24  25 26 27nltk.download('punkt')28nltk.download('stopwords')29nltk.download('wordnet')30 31def preprocess(idx,text,total_length):32    sentences = sent_tokenize(text)33    stop_words = set(stopwords.words('english'))34    lemmatizer = WordNetLemmatizer()35 36    padding=''+'0'*(len(str(total_length))-len(str(idx)))37    output=[]38    for sentence in sentences:39      output.append(' '.join([lemmatizer.lemmatize(word) for word in sentence.split() if word not in stop_words])+'@'+padding+str(idx))40    return output41 42def get_pre_processed_data(size):43  sentences=[]44  for idx,x in enumerate(df['Plot'].head(size).tolist()):45     sentences.extend(preprocess(idx,x,df.shape[0]))46  return sentences47 48#building_the_max_heap49def heapsort(np_array,k):50  h=[] 51  for idx,score in enumerate(np_array):52    heapq.heappush(h,(-score,idx))                    #max_heap53  return h54 55 56#return the id's of the movie57def get_top_k_matches(np_array,k,sentences):58   indices=set()59   h=heapsort(np_array,k)60 61   visited=set()62   indices=[]63   while h and len(indices)!=k:64     score,idx=heapq.heappop(h)65     i=len(sentences[idx])-1                  #based on the index find the sentence- reason for storing idx but not sentence66     count=167     number=068     while sentences[idx][i]!='@':                     #O(8-10 digits) i.e O(1) time69       number=number+count*int(sentences[idx][i])70       count*=1071       i-=172     73     if number not in visited:                #duplicate ids are not added, mainting 2 arrays is to maintian the order74       indices.append(number)75       visited.add(number)76   return indices77    78 79 80df=pd.DataFrame({"Title":dataset['train']['Title'],"Plot":dataset['train']['Overview']})81 82def getOutput(text, size=1000):83    sentences=get_pre_processed_data(int(size))84    np_array=getSimilarity(sentences,[text])85 86    output=[]87    for idx in get_top_k_matches(np_array,5,sentences):88        output.append("title = "+df.iloc[idx]['Title']+" "*5+" Plot = "+df.iloc[idx]['Plot'])89    return output90iface = gr.Interface(fn=getOutput, 91                     inputs=[gr.inputs.Textbox(label="Text")], 92                     outputs=[Label() for i in range(5)],93                     examples=[['After doing the list of experiments A mad scientist declares himself as the god'],["Three men fight for the girl's love"]]94                    )95iface.launch(debug=True)96