Update app.py

app.py

@@ -54,74 +54,37 @@ 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):
 
-#     # Split text into segments using split_text
-#     segments = split_text(text)
+    # Split text into segments using split_text
+    segments = split_text(text)
 
-#     # Initialize empty list for predictions
-#     predictions = []
+    # 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)
+    # 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)
+    # 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)
 
-#         with torch.no_grad():
-#             outputs = model(input_ids, attention_mask=attention_mask)
+        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
+        # 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
-#         })
+        # 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)
-
-  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)
-
-    with torch.no_grad():
-      outputs = model(input_ids, attention_mask=attention_mask)
-
-    probs, preds = extract_predictions(outputs)
-
-    predictions.append({
-      "segment_text": segment,
-      "label": model.config.id2label[preds[0]],  # assuming single label prediction
-      "probability": probs[preds[0]]
-    })
-
-  return predictions
-
-        
-
-# def classify_text(text):
-#   """
-#   This function preprocesses, feeds text to the model, and outputs the predicted class.
-#   """
-#   inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
-#   outputs = model(**inputs)
-#   logits = outputs.logits  # Access logits instead of pipeline output
-#   predictions = torch.argmax(logits, dim=-1)  # Apply argmax for prediction
-#   return model.config.id2label[predictions.item()]  # Map index to class label
 
 interface = gr.Interface(
     fn=classify_text,