import gradio as gr
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

# Load the model and tokenizer from Hugging Face Hub
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertForSequenceClassification.from_pretrained('jdmartinev/imdbreviews_classification_distilbert_v02')

# Define the function to perform text classification
def classify_text(input_text):
    inputs = tokenizer(input_text, return_tensors="pt")
    id2label = {0: "negative", 1: "positive"}
    # Get model predictions
    outputs = model(**inputs)
    pred = np.argmax(outputs)# Print logits
    return id2label[pred]
    

# Create Gradio interface
iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=2, placeholder="Enter text to classify"),
    outputs=gr.Tex(label="Class"),
    title="IMDB Review Classifier",
    description="Classify IMDB reviews using a fine-tuned DistilBERT model with LoRA.",
)

# Launch the app
iface.launch()