Create app.py

app.py

@@ -0,0 +1,41 @@
+# prompt: write a gradio app to infer the labels from the model we previously trained
+
+import gradio as gr
+import torch
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+# Load the fine-tuned model and tokenizer
+checkpoint = "25b3nk/ollama-issues-classifier"  # Replace with the actual path to your checkpoint directory
+model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
+tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+
+# Move the model to GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+
+# Function to perform inference
+def predict(text):
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True).to(device)
+    outputs = model(**inputs)
+    logits = outputs.logits
+    probabilities = torch.sigmoid(logits)  # Use sigmoid for multi-label classification
+
+    # Get predicted labels based on a threshold (e.g., 0.5)
+    predicted_labels = (probabilities > 0.5).nonzero()[:, 1].tolist()
+
+    # Map label IDs back to label names
+    predicted_labels_names = [model.config.id2label[label_id] for label_id in predicted_labels]
+
+    return predicted_labels_names
+
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict,
+    inputs=gr.Textbox(lines=5, placeholder="Enter the issue text here..."),
+    outputs=gr.Label(num_top_classes=len(model.config.id2label)),  # Display predicted labels
+    title="Issue Label Prediction",
+    description="Enter an issue description to predict its labels.",
+)
+
+iface.launch(debug=True)