Update README.md

README.md

@@ -62,14 +62,31 @@ Using HuggingFace Transformers (Python):
 --- 
 <small>
 
-from transformers import AutoTokenizer, AutoModelForSequenceClassification  
 import torch  
+from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification  
+from huggingface_hub import hf_hub_download  
+from safetensors.torch import load_file  
 
-model_id = "gameresearch/modernbert-eie-2048"  
-threshold = 0.2  # optimized decision threshold for EIE  
+base_model_id = "answerdotai/ModernBERT-base"  
+fine_tuned_model_id = "gameresearch/modernbert-eie-2048"  
+threshold = 0.2  # optimized decision threshold  
 
-tokenizer = AutoTokenizer.from_pretrained(model_id)  
-model = AutoModelForSequenceClassification.from_pretrained(model_id)  
+tokenizer = AutoTokenizer.from_pretrained(  
+    fine_tuned_model_id,  
+    subfolder="model",  
+)  
+
+config = AutoConfig.from_pretrained(base_model_id)  
+config.num_labels = 2  
+
+model = AutoModelForSequenceClassification.from_config(config)  
+
+weights_filename = "model/model.safetensors"  
+weights_path = hf_hub_download(fine_tuned_model_id, weights_filename)  
+
+state_dict = load_file(weights_path)  
+model.load_state_dict(state_dict)  
+model.eval()  
 
 def predict_eie(review_text: str):  
     inputs = tokenizer(  
@@ -77,7 +94,7 @@ def predict_eie(review_text: str):
         truncation=True,  
         padding="max_length",  
         max_length=2048,  
-        return_tensors="pt"  
+        return_tensors="pt",  
     )  
     with torch.no_grad():  
         outputs = model(**inputs)  
@@ -85,13 +102,9 @@ def predict_eie(review_text: str):
         probs = torch.softmax(logits, dim=-1)  
         prob_eie = probs[1].item()  
         label = int(prob_eie >= threshold)  
-    return {  
-        "prob_eie": prob_eie,  
-        "label": label  # 1 = EIE present, 0 = EIE absent  
-}  
+    return {"prob_eie": prob_eie, "label": label}  
 
-example = "This game absolutely destroyed me emotionally. I still think about the ending."  
-print(predict_eie(example))  
+print(predict_eie("This game absolutely destroyed me emotionally. I still think about the ending."))   
 </small>
 
 ---