Update app.py

app.py

@@ -1,22 +1,26 @@
 import os
 from fastapi import FastAPI
+from pydantic import BaseModel
 from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast
 import torch
 
-# Set Hugging Face cache directory to a writable path
 os.environ["HUGGINGFACE_HUB_CACHE"] = "/tmp/huggingface"
 
 app = FastAPI()
 
-model_name = "ADKU/ResearchGPT_model"  # Replace with your actual Hugging Face model ID
+model = DistilBertForSequenceClassification.from_pretrained("ADKU/ResearchGPT_model", cache_dir="/tmp/huggingface")
+tokenizer = DistilBertTokenizerFast.from_pretrained("ADKU/ResearchGPT_model", cache_dir="/tmp/huggingface")
 
-# Load model and tokenizer
-model = DistilBertForSequenceClassification.from_pretrained(model_name, cache_dir="/tmp/huggingface")
-tokenizer = DistilBertTokenizerFast.from_pretrained(model_name, cache_dir="/tmp/huggingface")
+class InputText(BaseModel):
+    inputs: str
 
 @app.post("/predict/")
-async def predict(text: str):
-    inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+async def predict(data: InputText):
+    inputs = tokenizer(data.inputs, return_tensors="pt", padding=True, truncation=True)
     outputs = model(**inputs)
     prediction = torch.argmax(outputs.logits, dim=-1).item()
     return {"prediction": prediction}
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0")