Update README.md

README.md

@@ -84,51 +84,50 @@ Two-class stress probability:
 
 ## How to Use
 
-Below is a self-contained example that dynamically downloads both the model code (`models.py`) and the weights from the Hub:
+Below is a self-contained example that dynamically downloads both the model code (`models.py`) and the weights from the Hub, then runs inference via the Hugging Face Transformers API—all in one script:
 
 ```python
 from huggingface_hub import hf_hub_download
 import importlib.util
+from transformers import AutoConfig, AutoModelForAudioClassification
 import torch
 import torch.nn.functional as F
 
-# 1) Download and load the model definition
-code_path = hf_hub_download(
-    repo_id="forwarder1121/voice-based-stress-recognition",
-    filename="models.py"
-)
-spec = importlib.util.spec_from_file_location("models", code_path)
-models = importlib.util.module_from_spec(spec)
-spec.loader.exec_module(models)
-StudentNet = models.StudentNet
-
-# 2) Download the checkpoint
-ckpt_path = hf_hub_download(
-    repo_id="forwarder1121/voice-based-stress-recognition",
-    filename="pytorch_model.bin"
-)
-
-# 3) Instantiate and load weights
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = StudentNet().to(device)
-
-state = torch.load(ckpt_path, map_location=device, weights_only=True)
-model.load_state_dict(state)
-model.eval()
-
-# ────── Create x_w2v here ──────
-# In real use, you would insert your preprocessed W2V embedding.
-# For testing, we create a random tensor of shape (1, 512).
-x_w2v = torch.randn(1, 512, device=device)
-
-# 4) Inference
-with torch.no_grad():
-    logits, _ = model(x_w2v)
-    probs     = F.softmax(logits, dim=1)
-    stress_p  = (probs[0, 1] * 100).item()  # Probability in [0, 100]
-
-print(f"Stress probability: {stress_p:.2f}%")
-
+def main():
+    repo = "forwarder1121/voice-based-stress-recognition"
+
+    # 1) Dynamically download & load the custom models.py
+    code_path = hf_hub_download(repo_id=repo, filename="models.py")
+    spec = importlib.util.spec_from_file_location("models", code_path)
+    models = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(models)
+    # now we have models.StudentForAudioClassification and models.StressConfig
+
+    # 2) Load config & model via Transformers (with remote code trust)
+    cfg = AutoConfig.from_pretrained(repo, trust_remote_code=True)
+    model = AutoModelForAudioClassification.from_pretrained(
+        repo,
+        trust_remote_code=True,
+        torch_dtype="auto"
+    )
+    model.eval()
+
+    # 3) Prepare a dummy W2V embedding for testing
+    #    In real use, replace this with your (1, 512) pre-computed W2V tensor.
+    batch_size = 1
+    DIM_W2V = 512
+    x_w2v = torch.randn(batch_size, DIM_W2V, dtype=next(model.parameters()).dtype)
+
+    # 4) Inference
+    with torch.no_grad():
+        outputs = model(x_w2v)                # SequenceClassifierOutput
+        probs   = F.softmax(outputs.logits, dim=-1)
+
+    print(f"Not stressed: {probs[0,0]*100:.1f}%")
+    print(f"Stressed    : {probs[0,1]*100:.1f}%")
+
+if __name__ == "__main__":
+    main()
 ```
 
 ## Citation