Update README.md

README.md

@@ -76,12 +76,11 @@ def zero_shot_eval(logger, device,
 
 if __name__ == "__main__":
     # get datasets
-    raw_datasets = load_dataset("Jiqing/ProtST-SubcellularLocalization", cache_dir="~/.cache/huggingface/datasets", split='test') # cache_dir defaults to "~/.cache/huggingface/datasets"
-    
-    #device = torch.device("cuda:0")
+    raw_datasets = load_dataset("mila-intel/ProtST-SubcellularLocalization", cache_dir="~/.cache/huggingface/datasets", split='test') # cache_dir defaults to "~/.cache/huggingface/datasets"
+
     device = torch.device("cpu")
     
-    protst_model = AutoModel.from_pretrained("Jiqing/ProtST-esm1b", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
+    protst_model = AutoModel.from_pretrained("mila-intel/ProtST-esm1b", trust_remote_code=True, torch_dtype=torch.bfloat16).to(device)
     protein_model = protst_model.protein_model
 +   import intel_extension_for_pytorch as ipex
 +   from optimum.intel.generation.modeling import jit_trace