Go back to Dr.Hassan's model + add question classify

models/space_classify.py

@@ -1,33 +1,55 @@
-# space_classify.py
-
 from transformers import pipeline
 
+
 def classifyA(text):
     """
     Extracts labels and scores from the input data,
     maps the labels using the provided mapping dictionary,
     and returns a list of formatted label-score strings.
     """
-    classification = pipeline(
-        task="text-classification",
-        model="asafaya/bert-base-arabic",
-        return_all_scores=True
-    )
-
+    from transformers import pipeline
+    classification = pipeline(task="text-classification", model="Hashuz/AS_MentalQAU", return_all_scores=True)
     result = []
-    # Define mapping dictionary
+
     mapping = {
-        'LABEL_0': 'تقديم معلومة',
-        'LABEL_1': 'توجيه أو ارشاد',
-        'LABEL_2': 'دعم نفسي'
+        'info': 'تقديم معلومة',
+        'guid': 'توجيه أو ارشاد',
+        'support': 'دعم نفسي'
     }
-
     output = classification(text)
-    for item in output[0]:  # assumes the structure [[{...}, {...}]]
+    for item in output[0]:  
         label = item['label']
         label = mapping.get(label)
         score = item['score']
         if score > 0.5:
             result.append(label)
+    return ', '.join(result)
 
+
+def classifyQ(text):
+    """
+    Extracts labels and scores from the input data,
+    maps the labels using the provided mapping dictionary,
+    and returns a list of formatted label-score strings.
+    """
+    from transformers import pipeline
+    classification = pipeline(task="text-classification", model="Hashuz/QT_MentalQA", return_all_scores=True)
+    result = []
+    
+    mapping = {
+        'diagnosis': 'فحص',
+        'treatment': 'علاج',
+        'anatomy': 'التشريح',
+        'epidemiology': 'الأوبئة',
+        'lifestyle': 'نمط الحياة',
+        'provider': 'مقدم الخدمة',
+        'other': 'غير محدد'
+    }
+    output = classification(text)
+    for item in output[0]: 
+        label = item['label']
+        label = mapping.get(label)
+        score = item['score']
+        if score > 0.5:
+            result.append(label)
     return ', '.join(result)