Added faster model switching and truncation to prevent errors on long inputs

app.py

@@ -1,4 +1,4 @@
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 import gradio as gr
 import torch
 
@@ -7,13 +7,18 @@ if torch.cuda.is_available():
 else:
     device = torch.device("cpu")
 
-summary = pipeline(task="summarization", model="facebook/bart-large-cnn", device=device)
-oracle = pipeline(task="zero-shot-classification", model="facebook/bart-large-mnli", device=device)
+summary_checkpoint = "facebook/bart-large-cnn" #"google/pegasus-large"
+oracle_checkpoint  = "facebook/bart-large-mnli"
+tokenizer = AutoTokenizer.from_pretrained(summary_checkpoint, device=device)
+summary = pipeline(task="summarization", model=summary_checkpoint, tokenizer=tokenizer, device=device)
+
+oracle = pipeline(task="zero-shot-classification", model=oracle_checkpoint, device=device)
 labels = ["merge","revert","fix","feature","update","refactor","test","security","documentation","style"]
+selected_labels = ["feature","update","refactor","test","security","documentation","style"]
 
 def do_the_thing(input, labels):
     #print(labels)
-    summarisation = summary(input)[0]['summary_text']
+    summarisation = summary(input, truncation=True)[0]['summary_text']
     zsc_results = oracle(sequences=[input, summarisation], candidate_labels=labels, multi_label=False, batch_size=2)
     classifications_input = {}
     for i in range(len(labels)):
@@ -32,7 +37,7 @@ with gr.Blocks() as frontend:
     btn_submit = gr.Button(value="Summarise and Classify")
     with gr.Row():
         with gr.Column():
-            input_labels = gr.Dropdown(label="Classification Labels", choices=labels, multiselect=True, value=labels, interactive=True, allow_custom_value=True, info="Labels to classify the original text and summary")
+            input_labels = gr.Dropdown(label="Classification Labels", choices=labels, multiselect=True, value=selected_labels, interactive=True, allow_custom_value=True, info="Labels to classify the original text and summary")
         with gr.Column():
             output_summary_text = gr.TextArea(label="Summary of Notes")
     with gr.Row():