numBery commited on
Commit
7f3cfe8
1 Parent(s): c636b3d

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +87 -0
app.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+
3
+ from keybert import KeyBERT
4
+ from keyphrase_vectorizers import KeyphraseCountVectorizer
5
+ from transformers import T5ForConditionalGeneration,T5Tokenizer
6
+
7
+ import nltk
8
+ from nltk.tokenize import sent_tokenize
9
+ nltk.download('stopwords')
10
+ nltk.download('punkt')
11
+
12
+ import streamlit as st
13
+
14
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
+
16
+ # Load KeyBert Model
17
+ kw_extractor = KeyBERT('all-MiniLM-L6-v2')
18
+ #kw_extractor = KeyBERT('distilbert-base-nli-mean-tokens')
19
+
20
+ # Load T5 for Paraphrasing
21
+ t5_model = T5ForConditionalGeneration.from_pretrained('ramsrigouthamg/t5_paraphraser')
22
+ t5_tokenizer = T5Tokenizer.from_pretrained('t5-base')
23
+ t5_model = t5_model.to(device)
24
+
25
+
26
+ doc = st.text_area("Enter a custom document")
27
+
28
+ def get_keybert_results_with_vectorizer(text, number_of_results=20):
29
+ keywords = kw_extractor.extract_keywords(text, vectorizer=KeyphraseCountVectorizer(), stop_words=None, top_n=number_of_results)
30
+ return keywords
31
+
32
+ def t5_paraphraser(text, number_of_results=10):
33
+ text = "paraphrase: " + text + " </s>"
34
+ max_len = 2048
35
+ encoding = t5_tokenizer.encode_plus(text, pad_to_max_length=True, return_tensors="pt")
36
+ input_ids, attention_masks = encoding["input_ids"].to(device), encoding["attention_mask"].to(device)
37
+
38
+ beam_outputs = t5_model.generate(
39
+ input_ids=input_ids, attention_mask=attention_masks,
40
+ do_sample=True,
41
+ max_length=2048,
42
+ top_k=50,
43
+ top_p=0.95,
44
+ early_stopping=True,
45
+ num_return_sequences=number_of_results
46
+ )
47
+
48
+ final_outputs =[]
49
+ for beam_output in beam_outputs:
50
+ sent = t5_tokenizer.decode(beam_output, skip_special_tokens=True, clean_up_tokenization_spaces=True)
51
+ final_outputs.append(sent)
52
+
53
+ return final_outputs
54
+
55
+
56
+ #### Extract Sentences with Keywords -> Paraphrase multiple versions -> Extract Keywords again
57
+
58
+ def extract_paraphrased_sentences(article):
59
+
60
+ original_keywords = [i[0] for i in get_keybert_results_with_vectorizer(article)]
61
+
62
+ article_sentences = sent_tokenize(article)
63
+ target_sentences = [sent for sent in article_sentences if any(kw[0] in sent for kw in original_keywords)]
64
+
65
+ start1 = time.time()
66
+ t5_paraphrasing_keywords = []
67
+
68
+ for sent in target_sentences:
69
+ ### T5
70
+ t5_paraphrased = t5_paraphraser(sent)
71
+ t5_keywords = [get_keybert_results_with_vectorizer(i) for i in t5_paraphrased]
72
+ t5_keywords = [word[0] for s in t5_keywords for word in s]
73
+
74
+ t5_paraphrasing_keywords.extend(t5_keywords)
75
+
76
+ print(f'T5 Approach2 PARAPHRASING RUNTIME: {time.time()-start1}\n')
77
+
78
+ print('T5 Keywords Extracted: \n{}\n\n'.format(t5_paraphrasing_keywords))
79
+ print('----------------------------')
80
+ print('T5 Unique New Keywords Extracted: \n{}\n\n'.format([i for i in set(t5_paraphrasing_keywords)
81
+ if i not in original_keywords]))
82
+
83
+ return t5_paraphrasing_keywords
84
+
85
+
86
+
87
+