Model can now take in X or X and Ts

app.py

@@ -13,10 +13,19 @@ with open(model_path, "rb") as f:
     model = pickle.load(f)
 
 class InputData(BaseModel):
-    X: list  
+    X: list
+    T0: list | None = None
+    T1: list | None = None
 
 @app.post("/model-effect")
 def predict(data: InputData):
     X = np.array(data.X)
-    effect = model.effect(X).tolist()
+
+    if data.T0 is not None and data.T1 is not None:
+        T0 = np.array(data.T0)
+        T1 = np.array(data.T1)
+        effect = model.effect(X=X, T0=T0, T1=T1).tolist()
+    else:
+        effect = model.effect(X).tolist()
+
     return {"effect": effect}