Spaces:
Running
Running
File size: 4,072 Bytes
d862c41 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
import pandas as pd
import numpy as np
from ast import literal_eval
import yake
import spacy
from sklearn.metrics.pairwise import cosine_similarity
from sentence_transformers import SentenceTransformer
import os
class ScriptMatcher:
def __init__(self, data_path = None, model_name='paraphrase-mpnet-base-v2',dataframe = None):
"""
Initialize the SeriesMatcher object.
Parameters:
data_path (str): Path to the dataset file.
model_name (str): Name of the sentence transformer model. Default is 'paraphrase-mpnet-base-v2'.
"""
if data_path is not None:
self.dataset = pd.read_csv(data_path)
if dataframe is not None:
self.dataset = dataframe
self.model = SentenceTransformer(model_name)
self.kw_extractor = yake.KeywordExtractor("en", n=1, dedupLim=0.9)
self.k_dataset = pd.read_csv('models/Similarity_K_Dataset/K_Dataset.csv')
self._ent_type = ["PERSON","NORP","FAC","ORG","GPE","LOC","PRODUCT","EVENT","WORK","ART","LAW",
"LANGUAGE","DATE","TIME","PERCENT","MONEY","QUANTITY","ORDINAL","CARDINAL"]
self.embeddings_synopsis_list = np.load("models/Similarity_K_Dataset/plot_embeddings.npy")
self.plot_embedding_list = np.load("models/Similarity_K_Dataset/synopsis_embeddings.npy")
try:
self.nlp = spacy.load("en_core_web_sm")
except:
print("Downloading spaCy NLP model...")
os.system(
"pip install https://huggingface.co/spacy/en_core_web_sm/resolve/main/en_core_web_sm-any-py3-none-any.whl")
self.nlp = spacy.load("en_core_web_sm")
def extract_keywords(self, text):
"""
Extract keywords from a given text using the YAKE keyword extraction algorithm.
Parameters:
text (str): Text from which to extract keywords.
Returns:
str: A string of extracted keywords joined by spaces.
"""
extracted_keywords = self.kw_extractor.extract_keywords(text)
return " ".join([keywords[0] for keywords in extracted_keywords if keywords[0] not in self._ent_type])
def preprocess_text(self, text):
"""
Process a given text to replace named entities and extract keywords.
Parameters:
text (str): The text to process.
Returns:
str: Processed text with named entities replaced and keywords extracted.
"""
doc = self.nlp(text)
replaced_text = text
for token in doc:
if token.ent_type_ != "MISC" and token.ent_type_ != "":
replaced_text = replaced_text.replace(token.text, f"<{token.ent_type_}>")
return self.extract_keywords(replaced_text)
def find_similar_series(self, new_synopsis, genres_keywords,k=5):
"""
Find series similar to a new synopsis.
Parameters:
new_synopsis (str): The synopsis to compare.
k (int): The number of similar series to return.
Returns:
pd.DataFrame: A dataframe of the closest series.
"""
processed_synopsis = self.preprocess_text(new_synopsis)
genre_keywords = " ".join(genres_keywords)
print(genre_keywords)
synopsis_sentence = genre_keywords + self.extract_keywords(processed_synopsis)
synopsis_embedding = self.model.encode([synopsis_sentence])
cosine_similarity_matrix = 0.75 * cosine_similarity(synopsis_embedding, self.embeddings_synopsis_list) + 0.25 * cosine_similarity(synopsis_embedding,self.plot_embedding_list)
top_k_indices = cosine_similarity_matrix.argsort()[0, -k:][::-1]
closest_series = self.k_dataset.iloc[top_k_indices]
# Add scores column
closest_series["Score"] = cosine_similarity_matrix[0, top_k_indices]
return closest_series[["Series", "Genre","Score"]].to_dict(orient='records')
|