Spaces:
Sleeping
Sleeping
File size: 3,361 Bytes
f84e8cd a1a9d9b f84e8cd 6650c1a f84e8cd 270a4af f84e8cd 6650c1a f84e8cd |
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 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 |
import torch
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
import gradio
mapper = {"India Today": 0, "NDTV": 1, "The Indian Express": 2, "The Times Of India": 3, "The Hindu": 4}
rev_mapper = {k:v for v,k in mapper.items()}
source_path = "rg089/bert_newspaper_source"
title_path = "rg089/t5-headline-generation"
summary_path = "rg089/distilbart-summarization"
device = "cuda" if torch.cuda.is_available() else "cpu"
source_model = AutoModelForSequenceClassification.from_pretrained(source_path).to(device)
source_tokenizer = AutoTokenizer.from_pretrained(source_path)
title_model = AutoModelForSeq2SeqLM.from_pretrained(title_path).to(device)
title_tokenizer = AutoTokenizer.from_pretrained(title_path)
summary_model = AutoModelForSeq2SeqLM.from_pretrained(summary_path).to(device)
summary_tokenizer = AutoTokenizer.from_pretrained(summary_path)
def generate(model, tokenizer, test_samples, prefix="", max_length=256):
model.eval()
with torch.no_grad():
if type(test_samples) == str:
test_samples = prefix + test_samples
else:
for i in range(len(test_samples)):
test_samples[i] = prefix + test_samples[i]
with tokenizer.as_target_tokenizer():
inputs = tokenizer(
test_samples,
truncation=True,
padding="max_length",
max_length=max_length,
return_tensors="pt")
input_ids = inputs.input_ids.to(device)
attention_mask = inputs.attention_mask.to(device)
outputs = model.generate(input_ids, attention_mask=attention_mask, num_beams=10, max_length=max_length) #, min_length=50)
output_str = tokenizer.batch_decode(outputs, skip_special_tokens=True)
return output_str[0]
def classify(model, tokenizer, content, title):
model.eval()
with torch.no_grad():
model_inputs = tokenizer(title, content, padding=True, truncation=True, return_tensors="pt").to(device)
outputs = model(**model_inputs)
logits = outputs.logits
selected = logits.argmax(dim=-1).cpu().tolist()
answers = [rev_mapper[sel] for sel in selected]
return answers[0]
def main(content, classify_source=False):
output = ""
title = generate(title_model, title_tokenizer, content, prefix="headline: ")
output += f"Title: {title}\n"
if classify_source:
source = classify(source_model, source_tokenizer, content, title)
output += f"Source: {source}\n\n"
else:
output += "\n"
summary = generate(summary_model, summary_tokenizer, content, prefix="")
output += f"Summary: {summary}"
return output
title = "News Helper: Generate Headlines, Summary and Classify the Newspaper Source!"
description = """
The current sources supported for classification are: The Times of India, The Indian Express, NDTV, The Hindu and India Today.
"""
placeholder = "Enter the content of the article here."
iface = gradio.Interface(fn=main, inputs=[gradio.inputs.Textbox(lines=10, placeholder=placeholder, label='Article Content:'),
gradio.inputs.Checkbox(default=True, label='Classify the Source:')], outputs="textbox", title=title,
description=description)
iface.launch() |