Spaces:
Running
Running
| import requests | |
| import time | |
| import gradio as gr | |
| 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}"} #This is a read only API key, do not share please :D. | |
| #This is visible just to experiment, I cannot share environment variables. | |
| def query(payload): #Function to use the API. | |
| response = requests.post(API_URL, headers=headers, json=payload) | |
| return response.json() #Return Json. | |
| # ########################################################## | |
| # Function to process the output and print classifications # | |
| ############################################################ | |
| def classify_output(item): | |
| label_dict = { | |
| 'LABEL_0': 'Cognition', | |
| 'LABEL_1': 'Affect', | |
| 'LABEL_2': 'Self', | |
| 'LABEL_3': 'Motivation', | |
| 'LABEL_4': 'Attention', | |
| 'LABEL_5': 'Overt_Behavior', | |
| 'LABEL_6': 'Context' | |
| } | |
| output = query({ #Try to query the endpoint. | |
| "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, | |
| }) | |
| # Store classifications in a list | |
| classifications = [] | |
| # Find the item with the highest score | |
| highest_score_item = max(output[0], key=lambda x: x['score']) | |
| for item in output[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'])) | |
| # Construct and print the classification message | |
| if classifications: | |
| classification_str = ', '.join([f"'{label}' ({score:.2f})" for label, score in classifications]) | |
| output1 = f"The item you provided is classified as: {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 | |
| ######################################### | |
| ######## RUN GRADIO APP ################# | |
| ######################################### | |
| demo = gr.Interface(fn=classify_output, inputs="text", outputs="text") | |
| demo.launch() | |