Update server.py

server.py

@@ -25,6 +25,8 @@ def root():
 
 @app.post("/predict_sentiment")
 def predict_sentiment(query: PredictRequest):
+    fhe_model = FHEModelServer("deployment/financial_rating")
+
     encrypted_encoding = base64.b64decode(query.encrypted_encoding)
     evaluation_key = base64.b64decode(query.evaluation_key)
     prediction = fhe_model.run(encrypted_encoding, evaluation_key)
@@ -32,4 +34,15 @@ def predict_sentiment(query: PredictRequest):
     # Encode base64 the prediction
     encoded_prediction = base64.b64encode(prediction).decode()
     return {"encrypted_prediction": encoded_prediction}
+    
+@app.post("/predict_legal")
+def predict_legal(query: PredictRequest):
+    fhe_legal_model = FHEModelServer("deployment/legal_rating")
+    
+    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}