diff --git a/flows/huggingface_message_classification.yaml b/flows/huggingface_message_classification.yaml new file mode 100644 index 0000000..51f695f --- /dev/null +++ b/flows/huggingface_message_classification.yaml @@ -0,0 +1,50 @@ +id: hugging_face +namespace: blueprint + +inputs: + - id: message_id + type: STRING + +tasks: + - id: retrieve_data + type: io.kestra.plugin.jdbc.postgresql.Query + sql: SELECT user_id, message FROM customer.message WHERE message_id = '{{ inputs.message_id }}' + fetchType: FETCH_ONE + + - id: classification + type: io.kestra.plugin.huggingface.Inference + model: facebook/bart-large-mnli + apiKey: "{{ secret('HUGGINGFACE_API_KEY') }}" + inputs: "{{ json(outputs.retrieve_data.row).message }}" + parameters: + candidate_labels: + - "support" + - "warranty" + - "upsell" + - "help" + + - id: insert_category + type: io.kestra.plugin.jdbc.postgresql.Query + sql: UPDATE customer.message SET category = '{{ json(outputs.classification.output).labels[0] }}' WHERE message_id = '{{ inputs.message_id }}' + + +pluginDefaults: + - type: io.kestra.plugin.jdbc.postgresql + values: + url: jdbc:postgresql://"{{secret('POSTGRES_HOST')}}" + username: "{{secret('POSTGRES_USERNAME')}}" + password: "{{secret('POSTGRES_PASSWORD')}}" + +extend: + title: Use HuggingFace Inference API to classify customer message + description: >- + This flow retrieve data from a Postgres database and use the HuggingFace Inference API to classify a customer message. + + Documentation about the HugginFace Inference API endpoint and how to create an API key can be find [here](https://huggingface.co/docs/api-inference/index). + tags: + - AI + - SQL + - API + ee: false + demo: false + meta_description: This flow retrieve data from a Postgres database and use the HuggingFace Inference API to classify a customer message. \ No newline at end of file