Update server.py

server.py

@@ -10,7 +10,7 @@ current_dir = Path(__file__).parent
 
 # Load the model
 fhe_model = FHEModelServer("deployment/financial_rating")
-
+fhe_legal_model = FHEModelServer("deployment/legal_rating")
 class PredictRequest(BaseModel):
     evaluation_key: str
     encrypted_encoding: str
@@ -29,6 +29,16 @@ def predict_sentiment(query: PredictRequest):
     evaluation_key = base64.b64decode(query.evaluation_key)
     prediction = fhe_model.run(encrypted_encoding, evaluation_key)
 
+    # Encode base64 the prediction
+    encoded_prediction = base64.b64encode(prediction).decode()
+    return {"encrypted_prediction": encoded_prediction}
+
+@app.post("/predict_legal")
+def predict_sentiment(query: PredictRequest):
+    encrypted_encoding = base64.b64decode(query.encrypted_encoding)
+    evaluation_key = base64.b64decode(query.evaluation_key)
+    prediction = fhe_legal_model.run(encrypted_encoding, evaluation_key)
+
     # Encode base64 the prediction
     encoded_prediction = base64.b64encode(prediction).decode()
     return {"encrypted_prediction": encoded_prediction}