Spaces:
Sleeping
Sleeping
| 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() |