MegaTronX's picture
Rename app.py to qloraapp.py
53517ec verified
raw
history blame contribute delete
538 Bytes
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()