data-silence commited on
Commit
4f3ff91
·
verified ·
1 Parent(s): f2c34e9

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -0
app.py ADDED
@@ -0,0 +1,48 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
4
+
5
+ # Загрузка модели и токенизатора
6
+ model_name = "data-silence/any-news-sum"
7
+ tokenizer = AutoTokenizer.from_pretrained(model_name)
8
+ model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
9
+
10
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
11
+ model.to(device)
12
+
13
+ def generate_summary_with_special_tokens(text, max_length=512):
14
+ inputs = tokenizer(text, return_tensors="pt", max_length=max_length, truncation=True).to(device)
15
+
16
+ outputs = model.generate(
17
+ **inputs,
18
+ max_length=max_length,
19
+ num_return_sequences=1,
20
+ no_repeat_ngram_size=4,
21
+ )
22
+
23
+ generated_text = tokenizer.decode(outputs[0], skip_special_tokens=False)
24
+
25
+ # Разделение на заголовок и резюме
26
+ parts = generated_text.split('<title_resume_sep>')
27
+ title = parts[0].replace("<pad> ", "").strip()
28
+ resume = parts[1].replace("</s>", "").strip() if len(parts) > 1 else ""
29
+
30
+ return title, resume
31
+
32
+ def summarize(text):
33
+ title, resume = generate_summary_with_special_tokens(text)
34
+ return title, resume
35
+
36
+ # Создание интерфейса Gradio
37
+ iface = gr.Interface(
38
+ fn=summarize,
39
+ inputs=gr.Textbox(lines=10, label="Введите текст новости"),
40
+ outputs=[
41
+ gr.Textbox(label="Сгенерированный заголовок"),
42
+ gr.Textbox(label="Сгенерированное резюме")
43
+ ],
44
+ title="Генератор заголовков и резюме новостей",
45
+ description="Введите текст новости, и модель создаст для неё заголовок и краткое резюме."
46
+ )
47
+
48
+ iface.launch()