from transformers import AutoModelForSequenceClassification, AutoTokenizer
import gradio as gr

model_name = "MegaTronX/Llama-3.2-1B-Instruct-Selectolax-QLoRA"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, from_tf=False)

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

demo = gr.Interface(
    fn=predict, 
    inputs="text", 
    outputs="label"
)

demo.launch()