Create app.py

app.py

@@ -0,0 +1,77 @@
+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"
+    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()