changed the classify text

app.py

@@ -54,38 +54,64 @@ def extract_predictions(outputs):
     
 # a function that classifies text
 
-def classify_text(text):
-    # Define labels
-    labels = ["depression", "anxiety", "bipolar disorder", "schizophrenia", "PTSD", "OCD", "ADHD", "autism", "eating disorder", "personality disorder", "phobia"]
+# def classify_text(text):
+#     # Define labels
+#     labels = ["depression", "anxiety", "bipolar disorder", "schizophrenia", "PTSD", "OCD", "ADHD", "autism", "eating disorder", "personality disorder", "phobia"]
+
+#     # Split text into segments using split_text
+#     segments = split_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")
+#         input_ids = inputs["input_ids"].to(device)
+#         attention_mask = inputs["attention_mask"].to(device)
 
-    # Split text into segments using split_text
-    segments = split_text(text)
+#         with torch.no_grad():
+#             outputs = model(input_ids, attention_mask=attention_mask)
+
+#         # Extract predictions for each segment
+#         probs, preds = extract_predictions(outputs)  # Define this function based on your model's output
+
+#         # 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
+#         })
+
+def classify_text(text):
+    
+ 
+  segments=split_text(text)
 
-    # Initialize empty list for predictions
-    predictions = []
+  predictions = []
+  for segment in segments:
+    inputs = tokenizer([segment], padding=True, return_tensors="pt")
+    input_ids = inputs["input_ids"].to(device)
+    attention_mask = inputs["attention_mask"].to(device)
 
-    # Move device to GPU if available
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = model.to(device)
+    with torch.no_grad():
+      outputs = model(input_ids, attention_mask=attention_mask)
 
-    # Loop through segments, process, and store predictions
-    for segment in segments:
-        inputs = tokenizer([segment], padding=True, return_tensors="pt")
-        input_ids = inputs["input_ids"].to(device)
-        attention_mask = inputs["attention_mask"].to(device)
+    probs, preds = extract_predictions(outputs)
 
-        with torch.no_grad():
-            outputs = model(input_ids, attention_mask=attention_mask)
+    predictions.append({
+      "segment_text": segment,
+      "label": model.config.id2label[preds[0]],  # assuming single label prediction
+      "probability": probs[preds[0]]
+    })
 
-        # Extract predictions for each segment
-        probs, preds = extract_predictions(outputs)  # Define this function based on your model's output
+  return predictions
 
-        # 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
-        })
+        
 
 # def classify_text(text):
 #   """