Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import AutoTokenizer, AutoModelForSeq2SeqLM | |
| import nltk | |
| import torch | |
| model_name = "Timur1984/t5-base-title-generation" # тут конечно всем надо поставить свой акк и имя модели | |
| print("Loading model...") | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
| nltk.download("punkt") | |
| nltk.download('punkt_tab') | |
| def generate_titles(text, num_titles, temperature): | |
| inputs = tokenizer( | |
| ["summarize: " + text], | |
| return_tensors="pt", | |
| truncation=True, | |
| max_length=512 | |
| ) | |
| outputs = model.generate( | |
| **inputs, | |
| do_sample=True, | |
| temperature=temperature, | |
| num_return_sequences=num_titles | |
| ) | |
| decoded = tokenizer.batch_decode(outputs, skip_special_tokens=True) | |
| titles = [nltk.sent_tokenize(t.strip())[0] for t in decoded] | |
| return "\n".join(titles) | |
| interface = gr.Interface( | |
| fn=generate_titles, | |
| inputs=[ | |
| gr.Textbox(lines=15, label="Article text"), | |
| gr.Slider(1,10,value=5,step=1,label="Number of titles"), | |
| gr.Slider(0.1,1.5,value=0.7,step=0.05,label="Temperature") | |
| ], | |
| outputs="text", | |
| title="Article Title Generator" | |
| ) | |
| interface.launch() |