Spaces:
Running
Running
| import gradio as gr | |
| from transformers import pipeline | |
| # Define model names | |
| models = { | |
| "ModernBERT Slop Classifier v1": "underscore2/modernbert_base_slop_classifier", | |
| "ModernBERT Slop Classifier v2": "underscore2/modernbert_base_slop_classifier_v2", | |
| "ModernBert Large Slop Classifier v3 (Best accuracy)": "underscore2/modernbert_large_slop_classifier_v3" | |
| } | |
| # Define the mapping for user-friendly labels | |
| # Note: Transformers pipelines often output 'LABEL_0', 'LABEL_1'. | |
| # We handle potential variations like just '0', '1'. | |
| label_map = { | |
| "LABEL_0": "Human (0)", | |
| "0": "Human (0)", | |
| "LABEL_1": "LLM (1)", | |
| "1": "LLM (1)" | |
| } | |
| large_v3 = pipeline("text-classification", model="underscore2/modernbert_large_slop_classifier_v3", top_k=None) | |
| # Function to load the selected model and classify text | |
| def classify_text(model_name, text): | |
| try: | |
| if models[model_name] != "underscore2/modernbert_large_slop_classifier_v3": | |
| classifier = pipeline("text-classification", model=models[model_name], top_k=None) | |
| else: | |
| classifier = large_v3 | |
| predictions = classifier(text) | |
| # Process predictions to use friendly labels | |
| processed_results = {} | |
| if predictions and isinstance(predictions, list) and predictions[0]: | |
| # predictions[0] should be a list of label dicts like [{'label': 'LABEL_1', 'score': 0.9...}, ...] | |
| for pred in predictions[0]: | |
| raw_label = pred["label"] | |
| score = pred["score"] | |
| # Use the map to get a friendly name, fallback to the raw label if not found | |
| friendly_label = label_map.get(raw_label, raw_label) | |
| processed_results[friendly_label] = score | |
| return processed_results | |
| except Exception as e: | |
| # Handle potential errors during model loading or inference | |
| print(f"Error: {e}") | |
| # Return an error message suitable for gr.Label | |
| return {"Error": f"Failed to process: {e}"} | |
| # Create the Gradio interface | |
| interface = gr.Interface( | |
| fn=classify_text, | |
| inputs=[ | |
| gr.Dropdown( | |
| list(models.keys()), | |
| label="Select Model", | |
| value=list(models.keys())[2] # Default model | |
| ), | |
| gr.Textbox( | |
| lines=2, | |
| placeholder="Enter text to classify", # Corrected placeholder | |
| value="This is an example sentence." # Changed example text | |
| ) | |
| ], | |
| # The gr.Label component works well for showing classification scores | |
| outputs=gr.Label(num_top_classes=2), # Show both classes explicitly | |
| title="ModernBERT Slop Classifier", | |
| description="Enter a sentence to see the slop and confidence scores", # Updated description | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| interface.launch() | |