Update app.py

app.py

@@ -1,28 +1,17 @@
 import gradio as gr
 from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
+from huggingface_hub import hf_hub_download
+import numpy as np
 
-# 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')
+id2label = {0: "negative", 1: "positive"}
 
-# Define the function to perform text classification
-def classify_text(input_text):
+def classify_text(input_text,id2label):
     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]
-    
+    pred = np.argmax(outputs)
+    return(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()
+demo = gr.Interface(classify_text, "text", "text")
+demo.launch()