Skip to content

Instantly share code, notes, and snippets.

@stephenlb
Last active April 9, 2024 07:04
Show Gist options
  • Save stephenlb/35449c8506a6a7e405ed73aac443c305 to your computer and use it in GitHub Desktop.
Save stephenlb/35449c8506a6a7e405ed73aac443c305 to your computer and use it in GitHub Desktop.
PubNub Functions HuggingFace - use any model from HuggingFace Serverless API inside PubNub Functions
//
// ** **
// ** 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();
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment