Views
No views yet
distilbert-base-cased and fine-tuned to predict 4 word classes: subjects, actions, objects, and descriptors.pip install onnxruntime transformers huggingface_hub numpy1import numpy as np
2import onnxruntime as ort
3from transformers import AutoTokenizer
4from huggingface_hub import hf_hub_download
5
6onnx_model_path = hf_hub_download(
7 "mbalabash/distilbert_subjects_actions_objects_descriptors", "model.onnx")
8tokenizer = AutoTokenizer.from_pretrained("distilbert-base-cased")
9id2label = {0: "ACTION", 1: "SUBJECT", 2: "OBJECT", 3: "DESCRIPTOR"}
10
11session = ort.InferenceSession(onnx_model_path)
12
13
14def classify_word(word):
15 inputs = tokenizer(word, return_tensors="np")
16 output = session.run(None, {
17 "input_ids": inputs["input_ids"].astype(np.int64),
18 "attention_mask": inputs["attention_mask"].astype(np.int64)
19 })
20
21 predicted_class_id = np.argmax(output[0])
22 return id2label[predicted_class_id]
23
24
25test_words = ["run", "teacher", "apple", "beautiful"]
26for word in test_words:
27 print(f"{word}: {classify_word(word)}")
28
29# OUTPUT ->
30# run: ACTION
31# teacher: SUBJECT
32# apple: OBJECT
33# beautiful: DESCRIPTORnpm install @xenova/transformers1import { pipeline } from "@xenova/transformers";
2
3const classifier = await pipeline("text-classification", "mbalabash/distilbert_subjects_actions_objects_descriptors");
4
5const testWords = ["run", "teacher", "apple", "beautiful"];
6
7for (const word of testWords) {
8 const result = await classifier(word);
9 console.log(`${word}: ${result[0].label}`);
10}
11
12// OUTPUT ->
13// run: ACTION
14// teacher: SUBJECT
15// apple: OBJECT
16// beautiful: DESCRIPTORPrediction Results:
--------------------------------------------------
run -> ACTION (confidence: 91.83%)
jump -> ACTION (confidence: 98.68%)
deploy -> ACTION (confidence: 99.90%)
considering -> ACTION (confidence: 99.85%)
training -> ACTION (confidence: 51.41%)
make -> ACTION (confidence: 99.71%)
teacher -> SUBJECT (confidence: 99.96%)
doctor -> SUBJECT (confidence: 99.95%)
woman -> SUBJECT (confidence: 99.96%)
viewer -> SUBJECT (confidence: 99.95%)
wizard -> SUBJECT (confidence: 99.97%)
pilot -> SUBJECT (confidence: 85.33%)
banana -> OBJECT (confidence: 99.69%)
laptop -> OBJECT (confidence: 99.95%)
dog -> OBJECT (confidence: 50.54%)
pencil -> OBJECT (confidence: 99.91%)
flower -> OBJECT (confidence: 56.07%)
car -> OBJECT (confidence: 82.67%)
beautiful -> DESCRIPTOR (confidence: 99.90%)
urgent -> DESCRIPTOR (confidence: 57.71%)
successful -> DESCRIPTOR (confidence: 99.94%)
frequently -> DESCRIPTOR (confidence: 99.95%)
strategic -> DESCRIPTOR (confidence: 99.42%)
faithful -> DESCRIPTOR (confidence: 99.91%)