Spaces:
Running on Zero
Running on Zero
File size: 1,431 Bytes
98d1e65 1717c0c 98d1e65 137f7ef 98d1e65 5bdeead 137f7ef 5bdeead 98d1e65 1717c0c 98d1e65 5bdeead 98d1e65 1717c0c 98d1e65 5bdeead 98d1e65 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | from pathlib import Path
from src.task.config import ModelConfig
from src.task.model import load_model, load_classifier_model
from src.infra.tokenizer import CharacterTokenizer
from src.task.inference.classify import TextClassifier
import pandas as pd
import numpy as np
from src.infra.dataset.util import build_text
# Runtime configs
# corpus = "data/tiny_shakespear.txt"
checkpoint = "checkpoints/mh_tiny_text_classifier.pt"
# CLASS_TOKENIZER_PATH = "ag_news_tokenizer.json"
CLASS_TOKENIZER_PATH = Path("checkpoints/ag_news_tokenizer.json")
tokenizer = CharacterTokenizer.load(
CLASS_TOKENIZER_PATH
)
CLASS_NAMES = [
"World",
"Sports",
"Business",
"Sci/Tech",
]
# Conf
model_config = ModelConfig()
model_config.number_of_heads = 2
model_config.context_length = 384
# Load model/checkpoint from file
# model = load_classifier_model(
# checkpoint_path=checkpoint,
# config=model_config,
# vocab_size=tokenizer.vocab_size,
# )
model = load_model(
checkpoint_path=checkpoint,
config=model_config,
vocab_size=tokenizer.vocab_size,
type="CLASSIFIER"
)
# generator
text_classifier = TextClassifier(
model=model,
tokenizer=tokenizer,
context_length=model_config.context_length,
)
# do generate
prompt_text="World cup 2030 to be in saudi arabia"
output = text_classifier.classify(
prompt_text
)
print(prompt_text)
print(CLASS_NAMES[output]) |