Update server.py

server.py

@@ -1,3 +1,4 @@
+"""Server that will listen for GET requests from the client."""
 from fastapi import FastAPI
 from joblib import load
 from concrete.ml.deployment import FHEModelServer
@@ -9,7 +10,6 @@ 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
@@ -29,17 +29,6 @@ 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("/legal_rating")
-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}