import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = AutoModelForSequenceClassification.from_pretrained(
    "scfengv/TVL_GeneralLayerClassifier",
    id2label = {0: "Cheer", 1: "Game", 2: "Broadcast", 3: "Chat"},
    label2id = {"Cheer": 0, "Game": 1, "Broadcast": 2, "Chat": 3}
    )
tokenizer = AutoTokenizer.from_pretrained("scfengv/TVL_GeneralLayerClassifier")

inputs = tokenizer("地震", return_tensors = "pt")

with torch.no_grad():
    logits = model(**inputs).logits

predicted_class_id = logits.argmax().item()
print(f"Predicted class: {predicted_class_id}")