artemnech/enrut5-base

This is mt5-base model google/mt5-base in which only Russian and English tokens are left

The model has been fine-tuned for several tasks:

• translation (opus100 dataset)
• dialog (daily dialog dataset)

How to use:

# !pip install transformers sentencepiece

from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, T5Tokenizer
import torch

model_name = 'artemnech/enrut5-base'

model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate(text, **kwargs):
    model.eval()
    inputs = tokenizer(text, return_tensors='pt').to(model.device)
    with torch.no_grad():
        hypotheses = model.generate(**inputs,  **kwargs)
    return tokenizer.decode(hypotheses[0], skip_special_tokens=True)

print(generate('translate ru-en: Я боюсь, что я не завершу доклад в ближайшее время.', num_beams=4,))
# I fear I'm not going to complete the report in the near future.

print(generate("translate en-ru: I'm afraid that I won't finish the report on time.", num_beams=4, max_length = 30))
# Я боюсь, что я не завершу доклад в ближайшее время.

print(generate('dialog: user1>>: Hello', num_beams=2))
# Hi

print(generate('dialog: user1>>: Hello user2>>: Hi user1>>: Would you like to drink something?', num_beams=2))
# I would like to drink a glass of wine.

from collections import deque

context =deque([], maxlen=6)
while True:
    text = input()
    text = 'user1>>: ' + text
    context.append(text)
    answ = generate('dialog: ' + ' '.join(context), num_beams=3, do_sample = True, temperature=1.5)
    context.append('user2>>: ' + answ)

    print('bot: ', answ)