Update README.md

README.md

@@ -34,14 +34,20 @@ pip install torch rdkit safetensors huggingface_hub numpy
 ## Usage
 
 ```python
-import torch
+import os, sys, torch
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
-from modeling_bioxmol import GatedGraphNeuralNetwork
-from featurizer import smiles2graph
 
 REPO = "arslanmasood/BioXMol"
 
+# Fetch the model definition + featurizer from the repo and make them importable
+for _f in ["modeling_bioxmol.py", "featurizer.py"]:
+    _code_dir = os.path.dirname(hf_hub_download(REPO, _f))
+sys.path.insert(0, _code_dir)
+
+from modeling_bioxmol import GatedGraphNeuralNetwork
+from featurizer import smiles2graph
+
 def load_encoder(seed: int = 0):
     model = GatedGraphNeuralNetwork(n_edge=1, in_dim=75, n_conv=6, fc_dims=[1024, 128])
     weights = hf_hub_download(REPO, f"bioxmol_soft_seed{seed}.safetensors")