File size: 1,295 Bytes
fdc43eb
 
 
 
 
d372f93
fdc43eb
 
 
 
 
9f38841
fdc43eb
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
ba5e197
fdc43eb
 
ba5e197
 
 
 
 
 
 
 
 
 
fdc43eb
ba5e197
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
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()