Update backend/app/model.py

backend/app/model.py

@@ -4,6 +4,7 @@ import os
 
 MODEL_NAME = "google/gemma-1.1-2b-it"
 SAVE_PATH = "./backend/app/checkpoints"
+LABEL_MAP = {0: "low risk", 1: "medium risk", 2: "high risk"}
 
 # Load tokenizer once
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -22,7 +23,7 @@ def predict(input_data: str):
     model = load_model()
     model.eval()
     with torch.no_grad():
-        inputs = tokenizer(input_data, return_tensors="pt")
+        inputs = tokenizer(input_data, return_tensors="pt", truncation=True, padding=True, max_length=256)
         outputs = model(**inputs)
         predicted_class = torch.argmax(outputs.logits, dim=1).item()
-        return predicted_class
+        return LABEL_MAP.get(predicted_class, "unknown")