Update README.md

README.md

@@ -54,6 +54,80 @@ tags = tag(article)
 print(tags)
 ```
 
+## Longer documents
+
+Assuming paragraphs are divided by: '\n\n'.
+
+```python
+from transformers import T5ForConditionalGeneration,T5Tokenizer
+import itertools
+import re 
+
+model = T5ForConditionalGeneration.from_pretrained("efederici/text2tags")
+tokenizer = T5Tokenizer.from_pretrained("efederici/text2tags")
+
+article = '''
+Da bambino era preoccupato che al mondo non ci fosse più nulla da scoprire. Ma i suoi stessi studi gli avrebbero dato torto: insieme a James Watson, nel 1953 Francis Crick strutturò il primo modello di DNA, la lunga sequenza di codici che identifica ogni essere vivente, rendendolo unico e diverso da tutti gli altri. 
+La scoperta gli valse il Nobel per la Medicina. È uscita in queste settimane per Codice la sua biografia, Francis Crick — Lo scopritore del DNA, scritta da Matt Ridley, che racconta vita e scienza dell'uomo che capì perché siamo fatti così.
+'''
+
+def words(text):
+    input_str = text
+    output_str = re.sub('[^A-Za-z0-9]+', ' ', input_str)
+    return output_str.split()
+
+def is_subset(text1, text2):
+    return all(tag in words(text1.lower()) for tag in text2.split())
+
+def cleaning(text, tags):
+    return [tag for tag in tags if is_subset(text, tag)]
+    
+def get_texts(self, text, max_len):
+    texts = list(filter(lambda x : x != '', text.split('\n\n')))
+    lengths = [len(tokenizer.encode(paragraph)) for paragraph in texts]
+    output = []
+    for i, par in enumerate(texts):
+        index = len(output)
+        if index > 0 and lengths[i] + len(tokenizer.encode(output[index-1])) <= max_len:
+            output[index-1] = "".join(output[index-1] + par)
+        else:
+            output.append(par)
+    return output
+        
+def get_tags(self, text, generate_kwargs):
+    input_text = 'summarize: ' + text.strip().replace('\n', ' ')
+    tokenized_text = tokenizer.encode(input_text, return_tensors="pt")
+    with torch.no_grad():
+        tags_ids = model.generate(tokenized_text, **generate_kwargs)
+    
+    output = []
+    for tags in tags_ids:
+        cleaned = cleaning(
+            text, 
+            list(set(tokenizer.decode(tags, skip_special_tokens=True).split(', ')))
+        )
+        output.append(cleaned)
+    
+    return list(set(itertools.chain(*output))) 
+        
+def tag(self, text, max_len, generate_kwargs):
+    texts = self.get_texts(text, max_len)
+    all_tags = [self.get_tags(text, generate_kwargs) for text in texts]
+    flatten_tags = itertools.chain(*all_tags)
+    return list(set(flatten_tags))
+
+params = {
+    "min_length": 0,
+    "max_length": 30,
+    "no_repeat_ngram_size": 2,
+    "num_beams": 4,
+    "early_stopping": True,
+    "num_return_sequences": 4,
+}
+tags = tag(article, 512, params)
+print(tags)
+```
+
 ### Overview
 
 - Model: T5 ([it5-small](https://huggingface.co/gsarti/it5-small))