Update README.md

README.md

@@ -14,18 +14,29 @@ widget:
   example_title: "Технологии"
 ---
 ## keyT5. Base (small) version
+Supported languages: ru
 
-[Large version](https://huggingface.co/0x7194633/keyt5-large)
+Github - [text2keywords](https://github.com/0x7o/text2keywords/edit/main/README.md)
 
-[Base version](https://huggingface.co/0x7194633/keyt5-base)
 
+[Pretraining Large version](https://huggingface.co/0x7194633/keyt5-large)
+|
+[Pretraining Base version](https://huggingface.co/0x7194633/keyt5-base)
+
+# Usage
 Example usage (the code returns a list with keywords. duplicates are possible):
+
+[![Try Model Training In Colab!](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/0x7o/text2keywords/blob/main/example/keyT5_use.ipynb)
+
+```
+pip install transformers sentencepiece
+```
+
 ```python
 from itertools import groupby
 import torch
 from transformers import T5ForConditionalGeneration, T5Tokenizer
-
-model_name = "0x7194633/keyt5-base"
+model_name = "0x7194633/keyt5-large" # or 0x7194633/keyt5-base
 tokenizer = T5Tokenizer.from_pretrained(model_name)
 model = T5ForConditionalGeneration.from_pretrained(model_name)
 
@@ -34,19 +45,20 @@ def generate(text, **kwargs):
     with torch.no_grad():
         hypotheses = model.generate(**inputs, num_beams=5, **kwargs)
     s = tokenizer.decode(hypotheses[0], skip_special_tokens=True)
-    s = s.replace('; ', ';').replace(' ;', ';').lower().split(';')
+    s = s.replace('; ', ';').replace(' ;', ';').lower().split(';')[:-1]
     s = [el for el, _ in groupby(s)]
     return s
 
-
 article = """Reuters сообщил об отмене 3,6 тыс. авиарейсов из-за «омикрона» и погоды
 Наибольшее число отмен авиарейсов 2 января пришлось на американские авиакомпании 
 SkyWest и Southwest, у каждой — более 400 отмененных рейсов. При этом среди 
 отмененных 2 января авиарейсов — более 2,1 тыс. рейсов в США. Также свыше 6400 
 рейсов были задержаны."""
 
-
 print(generate(article, top_p=1.0, max_length=64))  
+# ['авиаперевозки', 'отмена авиарейсов', 'отмена рейсов', 'отмена авиарейсов', 'отмена рейсов', 'отмена авиарейсов']
+```
+# Training
+Go to the training notebook and learn more about it:
 
-# ['лабораторная диагностика', 'панкреатит', 'профилактика', 'профилактика заболеваний', 'научно-популярное', 'биотехнологии', 'здоровье']
-```
+[![Try Model Training In Colab!](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/0x7o/text2keywords/blob/main/example/keyT5_train.ipynb)