Update README.md

README.md

@@ -9,14 +9,47 @@ pipeline_tag: text-classification
 
 This is the repo for Gen AI final project
 
+## Info
+Original code: https://github.com/hyunwoongko/transformer
+
+My version: https://github.com/Agaresd47/transformer_SAE
+
+
 
 ## Usage
 ```python
-from transformers import AutoModel, AutoTokenizer
+import torch
+from transformers import AutoModelForSequenceClassification, AutoTokenizer
+import torch.nn.functional as F
+
+# Load the model and tokenizer
+model = AutoModelForSequenceClassification.from_pretrained("agaresd/your-model-name")
+tokenizer = AutoTokenizer.from_pretrained("agaresd/your-model-name")
 
-model = AutoModel.from_pretrained("agaresd/GEN-AI-Final-project")
-tokenizer = AutoTokenizer.from_pretrained("agaresd/GEN-AI-Final-project")
+# Define the label mapping
+label_mapping = {
+    0: "no emotion",
+    1: "anger ",
+    2: "disgust ",
+    3: "fear ",
+    4: "Emotion: Happy",
+    5: "Emotion: Sad",
+    6: "Emotion: surprise"
+}
 
-inputs = tokenizer("Hello!", return_tensors="pt")
+# Input text
+input_text = "happy"
+
+# Tokenize and get model outputs
+inputs = tokenizer(input_text, return_tensors="pt")
 outputs = model(**inputs)
-print(outputs)
+
+# Get logits, apply softmax, and find the predicted class
+logits = outputs.logits
+probabilities = F.softmax(logits, dim=-1)
+predicted_class = torch.argmax(probabilities, dim=-1).item()
+
+# Map the predicted class to a word
+predicted_label = label_mapping[predicted_class]
+print(f"Input: {input_text}")
+print(f"Predicted Label: {predicted_label}")