iimran commited on
Commit
b9bfbd5
·
verified ·
1 Parent(s): 5c62a9f

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ from transformers import BartTokenizer, BartForConditionalGeneration
4
+
5
+ # Load the model and tokenizer from Hugging Face hub
6
+ model_name = "iimran/SAM-TheSummariserV2"
7
+ tokenizer = BartTokenizer.from_pretrained(model_name)
8
+ model = BartForConditionalGeneration.from_pretrained(model_name)
9
+ model.eval() # Set the model to evaluation mode
10
+
11
+ # Function to summarize the input text
12
+ def summarize(input_text):
13
+ # Tokenize the input text with truncation (adjust max_length as needed)
14
+ inputs = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=1024)
15
+
16
+ # Create global attention mask: assign global attention to the first token (required by LED)
17
+ global_attention_mask = torch.zeros(inputs["input_ids"].shape, dtype=torch.long)
18
+ global_attention_mask[:, 0] = 1
19
+
20
+ # Generate the summary using beam search (you can adjust parameters as needed)
21
+ summary_ids = model.generate(
22
+ inputs["input_ids"],
23
+ attention_mask=inputs["attention_mask"],
24
+ global_attention_mask=global_attention_mask,
25
+ max_length=512,
26
+ num_beams=4,
27
+ early_stopping=True,
28
+ )
29
+
30
+ # Decode the generated ids to a summary string, skipping special tokens
31
+ summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
32
+ return summary
33
+
34
+ # Create a Gradio interface with a title, description, submit button, and larger input text area
35
+ iface = gr.Interface(
36
+ fn=summarize, # Function that handles the summarization
37
+ inputs=gr.Textbox(
38
+ label="Enter Text to Summarize",
39
+ lines=10, # Make the input area larger by increasing the number of lines
40
+ placeholder="Paste or type the text you want to summarize here...",
41
+ ),
42
+ outputs=gr.Textbox(
43
+ label="Summary",
44
+ lines=5, # Adjust output area size (number of lines)
45
+ placeholder="Summary will appear here..."
46
+ ),
47
+ live=False, # Disable live updates, use the submit button instead
48
+ allow_flagging="never", # Optionally disable flagging
49
+ title="SAM - The Summariser", # Title of the page
50
+ description="SAM is a model which will help summarize large knowledge base articles into small summaries.", # Description of the model
51
+ )
52
+
53
+ # Launch the interface
54
+ iface.launch()