from transformers import AutoTokenizer, AutoModelForSequenceClassification

tokenizer = AutoTokenizer.from_pretrained("Kendamarron/fineweb-edu-classifier-ja")
model = AutoModelForSequenceClassification.from_pretrained("Kendamarron/fineweb-edu-classifier-ja")

def predict(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=-1).item()
    return predicted_class

text = "富士山は、日本で最も有名な山であり、日本全土にわたる広大な自然公園の一つです。高さは3,776メートルで、日本で最も高い山です。富士山は、東京都、静岡県、山梨県の3つの県にまたがっています。"
print(predict(text))
# >> 2