Add application file"

app.py

@@ -7,19 +7,32 @@ import os
 ## Loading the model: ##
 #######################
 api_key = os.environ.get("HF_API_KEY_INFERENCE")
-API_URL = "https://api-inference.huggingface.co/models/chernandezc/distilbert-base-uncased-finetuned-items-multi-label-21122023-AUGMENTED" #Api endpoint.
-headers = {"Authorization": f"Bearer {api_key}"} #Api Key. 
-                                                                            
 
-def query(payload): #Function to use the API.
-	response = requests.post(API_URL, headers=headers, json=payload)
+#7 dimensions:
+API_URL1 = "https://api-inference.huggingface.co/models/chernandezc/EEMM_7_categories_WB" #Api endpoint.
+headers = {"Authorization": f"Bearer {api_key}"} #Api Key, eqwual for both.
+
+
+def query1(payload): #Function to use the API.
+	response = requests.post(API_URL1, headers=headers, json=payload)
+	return response.json() #Return Json.
+
+
+#3 levels:
+API_URL2 = "https://api-inference.huggingface.co/models/chernandezc/EEMM_3_dimensions_1201" #Api endpoint.
+
+
+def query2(payload): #Function to use the API.
+	response = requests.post(API_URL2, headers=headers, json=payload)
 	return response.json() #Return Json.
 
 # ##########################################################
 # Function to process the output and print classifications #
 ############################################################
 def classify_output(item):
-    label_dict = {
+
+  #Dictionary for dimensions.
+    label_dict1 = {
         'LABEL_0': 'Cognition',
         'LABEL_1': 'Affect',
         'LABEL_2': 'Self',
@@ -29,38 +42,80 @@ def classify_output(item):
         'LABEL_6': 'Context'
     }
 
-    output = query({ #Try to query the endpoint.
+  #Dictionary for levels.
+    label_dict2 = {
+        'LABEL_0': 'Social',
+        'LABEL_1': 'Psychological',
+        'LABEL_2': 'Physical'
+    }
+
+
+    output1 = query1({ #Try to query the endpoint.
       "inputs": item,
       })
+    output2 = query2({ #Try to query the endpoint.
+      "inputs": item,
+      })
+
+    # Initial minimal delay
+    min_delay = 1  # seconds
+    #If model is idle wait and try again.
+    while 'error' in output1 or 'error' in output2:
+        if 'error' in output1:
+            time.sleep(min(output1.get("estimated_time", min_delay), min_delay))
+            output1 = query1({"inputs": item})
 
-    # If the model is loading, wait and try again.
-    while 'error' in output:
-        time.sleep(output["estimated_time"]) #Sleep the estimated time and try again.
-        output = query({
-          "inputs": item,
-          })
+        if 'error' in output2:
+            time.sleep(min(output2.get("estimated_time", min_delay), min_delay))
+            output2 = query2({"inputs": item})
 
     # Store classifications in a list
-    classifications = []
+    classifications1 = []
+    classifications2 = []
 
     # Find the item with the highest score
-    highest_score_item = max(output[0], key=lambda x: x['score'])
+    highest_score_item1 = max(output1[0], key=lambda x: x['score'])
+    highest_score_item2 = max(output2[0], key=lambda x: x['score'])
 
-    for item in output[0]:
+    for item in output1[0]:
         # Check if the score is greater than or equal to 0.5
         if item['score'] >= 0.5:
             # Append the category and score to the classifications list
-            classifications.append((label_dict[item['label']], item['score']))
-            
+            classifications1.append((label_dict1[item['label']], item['score']))
+
+    for item in output2[0]:
+        # Check if the score is greater than or equal to 0.5
+        if item['score'] >= 0.5:
+            # Append the category and score to the classifications list
+            classifications2.append((label_dict2[item['label']], item['score']))
+
     # Construct and print the classification message
-    if classifications:
-        classification_str = ', '.join([f"{label} ({score:.2f})" for label, score in classifications])
+    if (classifications1 and classifications2):
+        classification_str1 = ', '.join([f"{label} ({score:.2f})" for label, score in classifications1])
+        classification_str2 = ', '.join([f"{label} ({score:.2f})" for label, score in classifications2])
+
+        output_clas_and_lev = f"For dimensions: {classification_str1}\nFor levels: {classification_str2}"
+
+        return output_clas_and_lev
+
+    elif classifications1 and not classifications2:
+        classification_str1 = ', '.join([f"{label} ({score:.2f})" for label, score in classifications1])
+
+        output_clas_no_lev = f"For dimensions: {classification_str1}\nFor levels: No classifications with a score of 0.5 or higher were found.\nHowever, the highest probability was for: '{label_dict2[highest_score_item2['label']]}' ({round(highest_score_item2['score'],2)})\n Use this classification with caution due to uncertainty"
+
+        return output_clas_no_lev
+
+    elif classifications2 and not classifications1:
+        classification_str2 = ', '.join([f"{label} ({score:.2f})" for label, score in classifications2])
+
+        output_lev_no_clas = f"For levels: {classification_str2}\nFor dimensions: No classifications with a score of 0.5 or higher were found.\nHowever, the highest probability was for: '{label_dict1[highest_score_item1['label']]}' ({round(highest_score_item1['score'],2)}) \n  Use this classification with caution due to uncertainty"
+
+        return output_lev_no_clas
 
-        output1 = classification_str
-        return output1
     else:
-        output2 = f"No classifications with a score of 0.5 or higher were found. \n     However, the highest probability was for: '{label_dict[highest_score_item['label']]}' ({round(item['score'],2)}) \n     Use this classification with caution due to uncertainty"
-        return output2
+        output_lev_no_no = f"No classification with a score of 0.5 or higher were found for both levels and dimensions\nThe highest probability for dimensions was: '{label_dict1[highest_score_item1['label']]}' ({round(highest_score_item1['score'],2)}\nThe highest probability for level was: '{label_dict2[highest_score_item2['label']]}' ({round(highest_score_item2['score'],2)} \n  Use this classification with caution due to uncertainty"
+
+        return output_lev_no_no
     
     #########################################
     ######## RUN GRADIO APP #################