Updated README

README.md

@@ -27,47 +27,53 @@ tags:
 
 ## Custom Model Class
 
-This model uses a custom architecture implemented in `model.py`, specifically the class:
-```
-BERT_FFNN
-```
-If you want to load this model locally or fine-tune it further, make sure you have `model.py` in your working directory or import it correctly.
+This model uses a custom architecture defined inside the `emo_detector/` module:
+
+  - `emo_detector/configuration_bert_ffnn.py` → `BertFFNNConfig`
+  - `emo_detector/modeling_bert_ffnn.py` → `BERT_FFNN`
+
+To load or fine-tune this model, you must download the full repository (including the `emo_detector/` folder).
+The recommended way is to use `snapshot_download()` from Hugging Face Hub.
 
 ## Installation
 ```bash
-pip install torch transformers
+pip install torch transformers huggingface_hub
 ```
 
 ## Usage
 
 ```python
-from transformers import AutoTokenizer
+import sys
 import torch
-from model import BERT_FFNN 
+from transformers import AutoTokenizer
+from huggingface_hub import snapshot_download
+
+# Download entire repository
+repo_dir = snapshot_download("NeuralNest05/emo-detector")
+sys.path.append(repo_dir)
+
+# Import custom architecture + config
+from emo_detector.configuration_bert_ffnn import BertFFNNConfig
+from emo_detector.modeling_bert_ffnn import BERT_FFNN
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained("NeuralNest05/emo-detector")
 
-# Load model
-config = {
-    "bert_model_name": "microsoft/deberta-v3-base",
-    "hidden_dims": [192, 96],
-    "output_dim": 5,
-    "dropout": 0.2,
-    "pooling": "attention",
-    "freeze_bert": False,
-    "freeze_layers": 0,
-    "use_layer_norm": True
-}
-model = BERT_FFNN(**config)
-model_path = hf_hub_download(repo_id="NeuralNest05/emo-detector", filename="pytorch_model.bin")
-model.load_state_dict(torch.load(model_path, map_location=DEVICE))
+# Load model config and architecture
+config = BertFFNNConfig.from_pretrained("NeuralNest05/emo-detector")
+model = BERT_FFNN(config)
+
+# Load weights
+model.load_state_dict(torch.load(f"{repo_dir}/pytorch_model.bin", map_location=DEVICE))
 model.to(DEVICE)
 model.eval()
 
 # Example prediction
 texts = ["I am very happy today!", "This is scary..."]
-encodings = tokenizer(texts, truncation=True, padding=True, return_tensors="pt")
+encodings = tokenizer(texts, truncation=True, padding=True, return_tensors="pt").to(DEVICE)
+
 with torch.no_grad():
     logits = model(**encodings)
     probs = torch.sigmoid(logits)