Edit app.py

app.py

@@ -21,14 +21,6 @@ id2label = {
     6: 'schizophrenia'
 }
 
-# Inference function
-def predict(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-    with torch.no_grad():
-        outputs = model(**inputs)
-    pred_id = outputs.logits.argmax(dim=1).item()
-    return id2label[pred_id]
-
 def clean_text(text):
     text = str(text)
     text = text.lower()
@@ -39,6 +31,15 @@ def clean_text(text):
     text = re.sub(r'\s+', ' ', text).strip()
     return text
 
+# Inference function
+def predict(text):
+    text = clean_text(text)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    pred_id = outputs.logits.argmax(dim=1).item()
+    return id2label[pred_id]
+
 # Gradio UI
 demo = gr.Interface(
     fn=predict,