Update app.py

app.py

@@ -1,7 +1,10 @@
-import gradio as gr
+from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
+app = FastAPI()
+
 ROBERTA_MODEL = "Unknownaut/entity-level-framing-news-roberta"
 BERT_MODEL = "Unknownaut/entity-level-framing-news-bert"
 
@@ -12,9 +15,16 @@ _current_tokenizer = None
 _current_model_name = None
 
 
+class RequestData(BaseModel):
+    sentence: str
+    entity: str
+    model: str  # "RoBERTa" or "BERT"
+
+
 def load_model(model_choice):
     global _current_model, _current_tokenizer, _current_model_name
 
+    # reuse if already loaded
     if _current_model_name == model_choice:
         return _current_model, _current_tokenizer
 
@@ -38,12 +48,18 @@ def load_model(model_choice):
     return model, tokenizer
 
 
-def predict(sentence, entity, model_choice):
-    model, tokenizer = load_model(model_choice)
+@app.get("/")
+def health():
+    return {"status": "ok"}
+
+
+@app.post("/predict")
+def predict(data: RequestData):
+    model, tokenizer = load_model(data.model)
 
     inputs = tokenizer(
-        sentence,
-        entity,
+        data.sentence,
+        data.entity,
         return_tensors="pt",
         truncation=True,
         max_length=160
@@ -53,17 +69,4 @@ def predict(sentence, entity, model_choice):
         outputs = model(**inputs)
         pred = torch.argmax(outputs.logits, dim=1).item()
 
-    return labels[pred]
-
-
-demo = gr.Interface(
-    fn=predict,
-    inputs=[
-        gr.Textbox(),
-        gr.Textbox(),
-        gr.Radio(["RoBERTa", "BERT"])
-    ],
-    outputs="text"
-)
-
-demo.launch(enable_queue=False)
+    return {"label": labels[pred]}