Update README.md

README.md

@@ -58,6 +58,18 @@ pred = torch.argmax(outputs.logits, dim=1).item()
 
 print(pred)  # Map this index back to label if needed
 ```
+## 📦 Label Mapping
+
+The model outputs a numerical label corresponding to a prompt category. Below is the mapping between label IDs and their respective categories:
+
+0: brainstorming  
+1: classification  
+2: closed_qa  
+3: creative_writing  
+4: general_qa  
+5: information_extraction  
+6: open_qa  
+7: summarization  
 
 ## 🛠️ Training Details