Update README.md

README.md

@@ -41,7 +41,14 @@ pip install git+https://github.com/BiomedSciAI/biomed-multi-alignment.git
 
 A simple example for a task already supported by `ibm/biomed.omics.bl.sm.ma-ted-400m`:
 ```python
+import os
 
+from fuse.data.tokenizers.modular_tokenizer.op import ModularTokenizerOp
+from fuse.data.utils.collates import CollateDefault
+
+from mammal.examples.dti_bindingdb_kd.task import DtiBindingdbKdTask
+from mammal.keys import CLS_PRED, SCORES
+from mammal.model import Mammal
 
 # Load Model
 model = Mammal.from_pretrained("ibm/biomed.omics.bl.sm.ma-ted-400m.dti_bindingdb_pkd")
@@ -68,8 +75,8 @@ batch_dict = nn_model.forward_encoder_only([sample_dict])
 batch_dict = DtiBindingdbKdTask.process_model_output(
     batch_dict,
     scalars_preds_processed_key="model.out.dti_bindingdb_kd",
-    norm_y_mean=norm_y_mean,
-    norm_y_std=norm_y_std,
+    norm_y_mean=5.79384684128215,
+    norm_y_std=1.33808027428196,
 )
 ans = {
     "model.out.dti_bindingdb_kd": float(batch_dict["model.out.dti_bindingdb_kd"][0])