Update app.py

app.py

@@ -1,21 +1,66 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+import gc
 
-MODEL_NAME = "Unknownaut/entity-level-framing-news-roberta"
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+ROBERTA_MODEL = "Unknownaut/entity-level-framing-news-roberta"
+BERT_MODEL = "Unknownaut/entity-level-framing-news-bert"
 
 labels = ["Legitimate", "Aggressor", "Defensive", "Neutral"]
 
-def predict(sentence, entity):
+_current_model = None
+_current_tokenizer = None
+_current_model_name = None
+
+
+def unload_model():
+    global _current_model, _current_tokenizer
+
+    if _current_model is not None:
+        del _current_model
+        del _current_tokenizer
+        gc.collect()
+
+
+def load_model(model_key):
+    global _current_model, _current_tokenizer, _current_model_name
+
+    if _current_model_name == model_key:
+        return _current_model, _current_tokenizer
+
+    unload_model()
+
+    if model_key == "roberta":
+        tokenizer = AutoTokenizer.from_pretrained(ROBERTA_MODEL)
+        model = AutoModelForSequenceClassification.from_pretrained(ROBERTA_MODEL)
+
+    elif model_key == "bert":
+        tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL)
+        model = AutoModelForSequenceClassification.from_pretrained(BERT_MODEL)
+
+    else:
+        raise ValueError("Invalid model")
+
+    model.eval()
+
+    _current_model = model
+    _current_tokenizer = tokenizer
+    _current_model_name = model_key
+
+    return model, tokenizer
+
+
+def predict(sentence, entity, model_choice):
+    model_key = "roberta" if model_choice == "RoBERTa" else "bert"
+
+    model, tokenizer = load_model(model_key)
+
     inputs = tokenizer(
         sentence,
         entity,
         return_tensors="pt",
         truncation=True,
-        max_length=160
+        max_length=128
     )
 
     with torch.inference_mode():
@@ -24,9 +69,14 @@ def predict(sentence, entity):
 
     return labels[pred]
 
+
 demo = gr.Interface(
     fn=predict,
-    inputs=["text", "text"],
+    inputs=[
+        gr.Textbox(label="Sentence"),
+        gr.Textbox(label="Entity"),
+        gr.Radio(["RoBERTa", "BERT"], label="Model")
+    ],
     outputs="text"
 )