code-kern-ai / bricks

Open-source natural language enrichments at your fingertips.
Apache License 2.0
449 stars 23 forks source link

[MODULE] - distilBERT stock sentiment classifier #304

Open LeonardPuettmannKern opened 1 year ago

LeonardPuettmannKern commented 1 year ago

Please describe the module you would like to add to bricks I have fine-tuned a BERT model on a lot of stock news. The model can be used via the HuggingFace inference api.

Do you already have an implementation?

import requests
import json

API_KEY: str = "<API_KEY_GOES_HERE>"
ATTRIBUTE: str = "headline" # only text attributes

def stock_bert_classifier(record):
    """
    Uses toxic-bert via Hugging Face Inference API to classify toxicity in text.
    """
    api_token = API_KEY
    inputs = record[ATTRIBUTE].text
    headers = {"Authorization": f"Bearer {api_token}"}
    response = requests.post("https://api-inference.huggingface.co/models/KernAI/stockBERT", headers=headers, json={"inputs": inputs})
    json_response = response.json()
    result = [
        {item["label"]: item["score"] for item in entry}
        for entry in json_response
    ]
    return json.dumps(result[0])

Additional context https://huggingface.co/KernAI/stockBERT

LeonardPuettmannKern commented 1 year ago

implemented with #310

LeonardPuettmannKern commented 1 year ago

The output of the except statement needs to be changed, as it does not conform with the structure of refinery.