Edward Baker commited on
Commit
07cb201
1 Parent(s): d6b5ce1

updating with proper gradio

Browse files
Files changed (1) hide show
  1. app.py +49 -4
app.py CHANGED
@@ -1,7 +1,52 @@
 
 
 
 
 
 
1
  import gradio as gr
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
- def greet(name):
4
- return "Hello " + name + "!!"
 
 
 
 
 
5
 
6
- iface = gr.Interface(fn=greet, inputs="text", outputs="text")
7
- iface.launch()
 
1
+ from transformers import LEDTokenizer, LEDForConditionalGeneration
2
+ import torch
3
+ import re
4
+ tokenizer = LEDTokenizer.from_pretrained("patrickvonplaten/led-large-16384-pubmed")
5
+ model = LEDForConditionalGeneration.from_pretrained("patrickvonplaten/led-large-16384-pubmed").to("cuda").half()
6
+
7
  import gradio as gr
8
+ import os
9
+ import docx2txt
10
+
11
+
12
+ tokenizer = LEDTokenizer.from_pretrained("patrickvonplaten/led-large-16384-pubmed")
13
+ model = LEDForConditionalGeneration.from_pretrained("patrickvonplaten/led-large-16384-pubmed", return_dict_in_generate=True).to("cuda")
14
+
15
+
16
+
17
+ def summarize(text_file):
18
+ file_extension = os.path.splitext(text_file.name)[1]
19
+ if file_extension == ".txt":
20
+ # Load text from a txt file
21
+ with open(text_file.name, "r", encoding="utf-8") as f:
22
+ text = f.read()
23
+ elif file_extension == ".docx":
24
+ # Load text from a Word file
25
+ text = docx2txt.process(text_file.name)
26
+ else:
27
+ raise ValueError(f"Unsupported file type: {file_extension}")
28
+
29
+ input_ids = tokenizer(text, return_tensors="pt").input_ids.to("cuda")
30
+ global_attention_mask = torch.zeros_like(input_ids)
31
+ # set global_attention_mask on first token
32
+ global_attention_mask[:, 0] = 1
33
+
34
+ sequences = model.generate(input_ids, global_attention_mask=global_attention_mask).sequences
35
+
36
+ summary = tokenizer.batch_decode(sequences)[0]
37
+
38
+
39
+
40
+ return text, summary
41
+
42
+
43
 
44
+ iface = gr.Interface(
45
+ fn=summarize,
46
+ inputs=gr.inputs.File(label="Upload a txt file or a Word file for the input text"),
47
+ outputs=[gr.outputs.Textbox(label="Original text"), gr.outputs.Textbox(label="Summary")],
48
+ title="Academic Paper Summarization Demo",
49
+ description="Upload a txt file or a Word file for the input text. Get a summary generated by a small T5 model from Hugging Face.",
50
+ )
51
 
52
+ iface.launch()