Update README.md

README.md

@@ -59,22 +59,19 @@ A 10-fold cross-validation showed high reliability across different emotional di
 You can use the model and tokenizer as follows:
 
 ```python
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer
 import torch
 
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained("hplisiecki/polemo-intensity")
 
 # Load the model
-model = AutoModel.from_pretrained("hplisiecki/polemo-intensity")
-
-# Define emotion columns
-emotion_columns = ['Happiness', 'Sadness', 'Anger', 'Disgust', 'Fear', 'Pride', 'Valence', 'Arousal']
+model = Model.from_pretrained("hplisiecki/polemo-intensity")
 
 # Test the model with a sample input
 inputs = tokenizer("This is a test input.", return_tensors="pt")
-outputs = model(**inputs)
+outputs = model(inputs['input_ids'], inputs['attention_mask'])
 
 # Print out the emotion ratings
-for emotion, rating in zip(emotion_columns, outputs):
+for emotion, rating in zip(['Happiness', 'Sadness', 'Anger', 'Disgust', 'Fear', 'Pride', 'Valence', 'Arousal'], outputs):
     print(f"{emotion}: {rating.item()}")