MentalQA / models /space_classify.py
NLPAlhuzali's picture
Go back to Dr.Hassan's model + add question classify
01e1ed6 verified
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.
"""
from transformers import pipeline
classification = pipeline(task="text-classification", model="Hashuz/AS_MentalQAU", return_all_scores=True)
result = []
mapping = {
'info': 'تقديم معلومة',
'guid': 'توجيه أو ارشاد',
'support': 'دعم نفسي'
}
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)
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)