Spaces:
Running on Zero
Running on Zero
| 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]) |