added the id2labe

app.py

@@ -54,6 +54,11 @@ def extract_predictions(outputs):
     
 # a function that classifies text
 
+ # Move device to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+class_names=list(model.config.id2label.values())
+
 def classify_text(text):
 
     # Split text into segments using split_text
@@ -62,10 +67,7 @@ def classify_text(text):
     # Initialize empty list for predictions
     predictions = []
 
-    # Move device to GPU if available
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
-
+   
     # Loop through segments, process, and store predictions
     for segment in segments:
         inputs = tokenizer([segment], padding=True, return_tensors="pt")
@@ -77,13 +79,15 @@ def classify_text(text):
 
         # Extract predictions for each segment
         probs, preds = extract_predictions(outputs)  # Define this function based on your model's output
-
+        pred_label=class_names[preds[0].item()]
         # Append predictions for this segment
         predictions.append({
             "segment_text": segment,
-            "label": preds[0],  # Assuming single label prediction
-            "probability": probs[preds[0]]  # Access probability for the predicted label
+            "label": pred_label,  # Assuming single label prediction
+            "probability": probs[0][preds[0]].item()  # Access probability for the predicted label
         })
+    
+    return predictions
 
 
 interface = gr.Interface(