Update app.py

app.py

@@ -1,7 +1,6 @@
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
-import gc
 
 ROBERTA_MODEL = "Unknownaut/entity-level-framing-news-roberta"
 BERT_MODEL = "Unknownaut/entity-level-framing-news-bert"
@@ -13,28 +12,17 @@ _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):
+def load_model(model_choice):
     global _current_model, _current_tokenizer, _current_model_name
 
-    if _current_model_name == model_key:
+    if _current_model_name == model_choice:
         return _current_model, _current_tokenizer
 
-    unload_model()
-
-    if model_key == "roberta":
+    if model_choice == "RoBERTa":
         tokenizer = AutoTokenizer.from_pretrained(ROBERTA_MODEL)
         model = AutoModelForSequenceClassification.from_pretrained(ROBERTA_MODEL)
 
-    elif model_key == "bert":
+    elif model_choice == "BERT":
         tokenizer = AutoTokenizer.from_pretrained(BERT_MODEL)
         model = AutoModelForSequenceClassification.from_pretrained(BERT_MODEL)
 
@@ -45,22 +33,20 @@ def load_model(model_key):
 
     _current_model = model
     _current_tokenizer = tokenizer
-    _current_model_name = model_key
+    _current_model_name = model_choice
 
     return model, tokenizer
 
 
 def predict(sentence, entity, model_choice):
-    model_key = "roberta" if model_choice == "RoBERTa" else "bert"
-
-    model, tokenizer = load_model(model_key)
+    model, tokenizer = load_model(model_choice)
 
     inputs = tokenizer(
         sentence,
         entity,
         return_tensors="pt",
         truncation=True,
-        max_length=160
+        max_length=128
     )
 
     with torch.inference_mode():
@@ -73,11 +59,12 @@ def predict(sentence, entity, model_choice):
 demo = gr.Interface(
     fn=predict,
     inputs=[
-        gr.Textbox(label="Sentence"),
-        gr.Textbox(label="Entity"),
-        gr.Radio(["RoBERTa", "BERT"], label="Model")
+        gr.Textbox(),
+        gr.Textbox(),
+        gr.Radio(["RoBERTa", "BERT"])
     ],
     outputs="text"
 )
 
-demo.launch()
+# 🔥 IMPORTANT: enable API
+app = gr.mount_gradio_app(None, demo, path="/")