News headline generation using neural nets, trained on "Rossiya Segodnya" news dataset.
Using last 10% of the "Rossiya Segodnya" news dataset for now. May switch to another dataset later.
You may want to create a virtual environment first
pip install -r requirements.txtBuild docker image
docker build . --tag headlines:latestfrom transformers import AutoTokenizer, AutoModelForSeq2SeqLM
MODEL_NAME = "dmitry-vorobiev/rubert_ria_headlines"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
text = "Скопируйте текст статьи / новости"
encoded_batch = tokenizer.prepare_seq2seq_batch(
[text],
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=512)
output_ids = model.generate(
input_ids=encoded_batch["input_ids"],
max_length=32,
no_repeat_ngram_size=3,
num_beams=5,
top_k=0
)
headline = tokenizer.decode(output_ids[0],
skip_special_tokens=True,
clean_up_tokenization_spaces=False)
print(headline)Run docker container interactive session:
docker run -it --name headlines_gpu --gpus all headlines:latestExecute evaluation script:
sh /app/scripts/eval_docker.shBy default, results can be found at /io/output.
-
2021-02-01 | huggingface | kaggle | TPUv3, batch size of 192, better eval results | Evaluation: last 1% of ria, last 10% of ria
-
2021-01-31 | huggingface | kaggle | same as previous, but trained on TPUv3 with higher batch size (256 vs 56). | Evaluation: last 1% of ria, last 10% of ria
-
2021-01-28 | huggingface | bert2bert, initialized with the
DeepPavlov/rubert-base-casedpretrained weights and fine-tuned on the first 90% of "Rossiya Segodnya" news dataset for 1.6 epochs.