File size: 2,550 Bytes
db0f499
 
 
8ed0852
 
f4f3763
63a8066
 
8ed0852
db0f499
6a0b78f
 
 
 
f4f3763
 
 
 
6a0b78f
 
4dfc57d
 
6a0b78f
 
 
09e7c03
6a0b78f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
09e7c03
db0f499
 
 
 
8ed0852
 
 
 
63a8066
8ed0852
 
 
f4f3763
6a0b78f
8ed0852
 
6a0b78f
 
f4f3763
e9043bf
db0f499
 
 
6a0b78f
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
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)"
}

# Function to load the selected model and classify text
def classify_text(model_name, text):
    try:
        classifier = pipeline("text-classification", model=models[model_name], top_k=None)
        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=models[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()