PubNub Functions HuggingFace - use any model from HuggingFace Serverless API inside PubNub Functions

huggingface-pubnub-function.js

//
// **                                              **
// **  Add your API Key to MY SECRETS (Left Panel) **
// **                                              **
// **    HUGGINGFACE_API_KEY as the Secret Name    **
// **                                              **
//
let HUGGINGFACE_API_KEY = null;
function getAPIKey() {
    // Use cached key
    if (HUGGINGFACE_API_KEY) {
        return new Promise(resolve => resolve(HUGGINGFACE_API_KEY));
    }
    // Fetch key from vault
    return vault.get("HUGGINGFACE_API_KEY").then(apikey => {
        HUGGINGFACE_API_KEY = apikey;
        return new Promise(resolve => resolve(HUGGINGFACE_API_KEY));
    });
}

//
// Import Modules
//
const http = require('xhr');
const vault = require('vault');

// 
// Main
// 
export default async (request) => {
    let message = request.message.text;

    // Sentiment
    let model = 'distilbert-base-uncased-finetuned-sst-2-english';
    let response = await query(message, model);
    console.log(response);

    // GPT2
    model = 'gpt2';
    response = await query(message, model);
    console.log(response);

    // Google's Gemma
    model = 'google/gemma-7b-it';
    response = await query(`<start_of_turn>user\n${message}\n<start_of_turn>model`, model);
    console.log(response);

    // Capture
    model = 'dbmdz/bert-large-cased-finetuned-conll03-english';
    response = await query(message, model);
    console.log(response);
    
    return request.ok()
};

// 
// HuggingFace API Call
// 
async function query(text, model='distilbert-base-uncased-finetuned-sst-2-english') {
    const apikey = await getAPIKey();
    const response = await http.fetch( `https://api-inference.huggingface.co/models/${model}`, {
        headers: { Authorization: `Bearer ${apikey}`, 'Content-Type': 'application/json' },
        method: "POST",
        body: JSON.stringify({inputs:text}),
    });
    return await response.json();
}