Spaces:
Runtime error
Runtime error
| import re | |
| import numpy as np | |
| import itertools | |
| import torch | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| from sentence_transformers import SentenceTransformer | |
| from sklearn.feature_extraction.text import CountVectorizer | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| class KeywordExtraction: | |
| def __init__(self, n_gram_range=(1, 1), stop_words='english', model_name='distilbert-base-nli-mean-tokens'): | |
| self.n_gram_range = n_gram_range | |
| self.stop_words = stop_words | |
| self.model_name = model_name | |
| self.model = SentenceTransformer(self.model_name) | |
| def __call__(self, doc, top_n=5, diversity=('mmr', 0.7)): | |
| doc_embedding = self.get_document_embeddings(doc) | |
| candidates = self.get_candidates(doc) | |
| candidate_embeddings = self.get_candidate_embeddings(candidates) | |
| try: | |
| if diversity[0] == 'mmr': | |
| # print('using maximal marginal relevance method...') | |
| return self.maximal_marginal_relevance(doc_embedding, | |
| candidate_embeddings, | |
| candidates, | |
| top_n=top_n, | |
| diversity=diversity[1]) | |
| elif diversity[0] == 'mss': | |
| # print('using max sum similarity method...') | |
| return self.max_sum_similarity(doc_embedding, | |
| candidate_embeddings, | |
| candidates, | |
| top_n=top_n, | |
| nr_candidates=diversity[1]) | |
| else: | |
| # print('using default method...') | |
| return self.get_keywords(doc_embedding, candidate_embeddings, candidates, top_n) | |
| except Exception as e: | |
| print(e) | |
| def get_candidates(self, doc): | |
| # Extract candidate words/phrases | |
| count = CountVectorizer(ngram_range=self.n_gram_range, stop_words=self.stop_words).fit([doc]) | |
| return count.get_feature_names_out() | |
| def get_candidate_embeddings(self, candidates): | |
| return self.model.encode(candidates) | |
| def get_document_embeddings(self, doc): | |
| return self.model.encode([doc]) | |
| def get_keywords(self, doc_embedding, candidate_embeddings, candidates, top_n=5): | |
| distances = cosine_similarity(doc_embedding, candidate_embeddings) | |
| keywords = [candidates[index] for index in distances.argsort()[0][-top_n:]] | |
| return keywords | |
| def max_sum_similarity(self, doc_embedding, candidate_embeddings, candidates, top_n, nr_candidates): | |
| # Calculate distances and extract keywords | |
| distances = cosine_similarity(doc_embedding, candidate_embeddings) | |
| distances_candidates = cosine_similarity(candidate_embeddings, | |
| candidate_embeddings) | |
| # Get top_n words as candidates based on cosine similarity | |
| words_idx = list(distances.argsort()[0][-nr_candidates:]) | |
| words_vals = [candidates[index] for index in words_idx] | |
| distances_candidates = distances_candidates[np.ix_(words_idx, words_idx)] | |
| # Calculate the combination of words that are the least similar to each other | |
| min_sim = np.inf | |
| candidate = None | |
| for combination in itertools.combinations(range(len(words_idx)), top_n): | |
| sim = sum([distances_candidates[i][j] for i in combination for j in combination if i != j]) | |
| if sim < min_sim: | |
| candidate = combination | |
| min_sim = sim | |
| return [words_vals[idx] for idx in candidate] | |
| def maximal_marginal_relevance(self, doc_embedding, word_embeddings, words, top_n, diversity): | |
| # Extract similarity within words, and between words and the document | |
| word_doc_similarity = cosine_similarity(word_embeddings, doc_embedding) | |
| word_similarity = cosine_similarity(word_embeddings) | |
| # Initialize candidates and already choose best keyword/keyphras | |
| keywords_idx = [np.argmax(word_doc_similarity)] | |
| candidates_idx = [i for i in range(len(words)) if i != keywords_idx[0]] | |
| for _ in range(top_n - 1): | |
| # Extract similarities within candidates and | |
| # between candidates and selected keywords/phrases | |
| candidate_similarities = word_doc_similarity[candidates_idx, :] | |
| target_similarities = np.max(word_similarity[candidates_idx][:, keywords_idx], axis=1) | |
| # Calculate MMR | |
| mmr = (1-diversity) * candidate_similarities - diversity * target_similarities.reshape(-1, 1) | |
| mmr_idx = candidates_idx[np.argmax(mmr)] | |
| # Update keywords & candidates | |
| keywords_idx.append(mmr_idx) | |
| candidates_idx.remove(mmr_idx) | |
| return [words[idx] for idx in keywords_idx] | |
| def regex(phrase, m=0, n=3): | |
| strng = "([\s]*[a-zA-Z0-9]*[\s]*){%d,%d}" % (m,n) | |
| return strng.join(phrase.split()) | |
| def remove_square_brackets(text): | |
| return re.sub('\[[0-9]+\]', '', text) | |
| def remove_extra_spaces(text): | |
| return re.sub('[\s]{2,}', ' ', text) | |
| def preprocess_text(text): | |
| text = re.sub('\[[0-9]+\]', '', text) | |
| text = re.sub('[\s]{2,}', ' ', text) | |
| text = text.strip() | |
| return text | |
| def sent_tokenize(text): | |
| sents = text.split('.') | |
| sents = [s.strip() for s in sents if len(s)>0] | |
| return sents | |
| def get_key_sentences(text, top_n=5, diversity=('mmr', 0.6)): | |
| kw_extractor = KeywordExtraction(n_gram_range=(1,3)) | |
| text = preprocess_text(text) | |
| sentences = sent_tokenize(text) | |
| key_phrases = kw_extractor(text, top_n=top_n, diversity=diversity) | |
| if key_phrases is None: | |
| return None | |
| key_sents = dict() | |
| for phrase in key_phrases: | |
| found = False | |
| for i, sent in enumerate(sentences): | |
| if re.search(regex(phrase), sent): | |
| found = True | |
| if i not in key_sents: | |
| key_sents[i] = sent | |
| if not found: | |
| print(f'The phrase "{phrase}" was not matched!') | |
| return key_sents | |
| class ParaphraseModel: | |
| def __init__(self, model_name="Vamsi/T5_Paraphrase_Paws"): | |
| self.model_name = model_name | |
| self.tokenizer = AutoTokenizer.from_pretrained(self.model_name) | |
| self.model = AutoModelForSeq2SeqLM.from_pretrained(self.model_name) | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| def __call__(self, inputs, top_k=200, top_p=0.95, num_sequences=5): | |
| text = self.prepare_list_input(inputs) if type(inputs) == type([]) else f"paraphrase: {inputs} </s>" | |
| encoding = self.tokenizer.batch_encode_plus(text, pad_to_max_length=True, return_tensors="pt") | |
| input_ids = encoding["input_ids"].to(self.device) | |
| attention_masks = encoding["attention_mask"].to(self.device) | |
| outputs = self.model.generate( | |
| input_ids=input_ids, attention_mask=attention_masks, | |
| max_length=256, | |
| do_sample=True, | |
| top_k=top_k, | |
| top_p=top_p, | |
| early_stopping=True, | |
| num_return_sequences=num_sequences | |
| ) | |
| lines = [] | |
| for output in outputs: | |
| line = self.tokenizer.decode(output, | |
| skip_special_tokens=True, | |
| clean_up_tokenization_spaces=True) | |
| lines.append(line) | |
| return lines | |
| def prepare_list_input(self, lst): | |
| sentences = [] | |
| for sent in lst: | |
| sentences.append(f"paraphrase: {sent} </s>") | |
| return sentences | |