Spaces:

rundac-inference
/

assistant-nickname-classifier

Paused

Update app.py

d31b99d verified 4 months ago

2.29 kB

import os
import torch
import gradio as gr
from transformers import pipeline

access_token = os.environ.get('hf_token')

if torch.cuda.is_available():
    device_arg = 0  # nvidia gpu
elif torch.backends.mps.is_available():
    device_arg = "mps"  # apple silicon gpu
else:
    device_arg = -1  # cpu fallback

# Initialize text classification pipeline
classifier = pipeline(
    task="text-classification",
    model="Falconsai/offensive_speech_detection",
    device=device_arg,
    return_all_scores=True  # Asks pipeline for both labels
)

# Human-readable mapping for model output
LABEL_MAPPING = {
    "LABEL_0": "Normal Speech",
    "LABEL_1": "Offensive Speech"
}

def detect_offensive_speech(text):
    if not text.strip():
        return {"Empty Input": 1.0}
    
    raw_output = classifier(text)
    
    # --- DEFENSIVE PARSING BLOCK ---
    # Case 1: Nested list [[{...}, {...}]] -> Extract inner list
    if isinstance(raw_output, list) and len(raw_output) > 0 and isinstance(raw_output[0], list):
        items = raw_output[0]
    # Case 2: Flat list [{...}, {...}] -> Use as-is
    elif isinstance(raw_output, list):
        items = raw_output
    # Case 3: Single dictionary {...} -> Wrap in a list
    elif isinstance(raw_output, dict):
        items = [raw_output]
    else:
        return {"Parsing Error": 1.0}
    # -------------------------------

    formatted_results = {}
    for item in items:
        # Extra safety: Ensure 'item' is a dictionary before reading keys
        if isinstance(item, dict) and "label" in item and "score" in item:
            clean_label = LABEL_MAPPING.get(item["label"], item["label"])
            formatted_results[clean_label] = float(item["score"])
            
    # Fallback if parsing completely fails to find labels
    if not formatted_results:
        return {"Unknown Error": 1.0}
        
    return formatted_results

# Define Gradio Interface
demo = gr.Interface(
    fn=detect_offensive_speech,
    inputs=gr.Textbox(placeholder="Enter text here...", lines=3),
    outputs=gr.Label(label="Classification Results", num_top_classes=2),
    title="Offensive Speech Detection",
    description="A local, optimized text-classification pipeline running on Hugging Face weights."
)

if __name__ == "__main__":
    demo.launch()

	import os
	import torch
	import gradio as gr
	from transformers import pipeline

	access_token = os.environ.get('hf_token')

	if torch.cuda.is_available():
	device_arg = 0 # nvidia gpu
	elif torch.backends.mps.is_available():
	device_arg = "mps" # apple silicon gpu
	else:
	device_arg = -1 # cpu fallback

	# Initialize text classification pipeline
	classifier = pipeline(
	task="text-classification",
	model="Falconsai/offensive_speech_detection",
	device=device_arg,
	return_all_scores=True # Asks pipeline for both labels
	)

	# Human-readable mapping for model output
	LABEL_MAPPING = {
	"LABEL_0": "Normal Speech",
	"LABEL_1": "Offensive Speech"
	}

	def detect_offensive_speech(text):
	if not text.strip():
	return {"Empty Input": 1.0}

	raw_output = classifier(text)

	# --- DEFENSIVE PARSING BLOCK ---
	# Case 1: Nested list [[{...}, {...}]] -> Extract inner list
	if isinstance(raw_output, list) and len(raw_output) > 0 and isinstance(raw_output[0], list):
	items = raw_output[0]
	# Case 2: Flat list [{...}, {...}] -> Use as-is
	elif isinstance(raw_output, list):
	items = raw_output
	# Case 3: Single dictionary {...} -> Wrap in a list
	elif isinstance(raw_output, dict):
	items = [raw_output]
	else:
	return {"Parsing Error": 1.0}
	# -------------------------------

	formatted_results = {}
	for item in items:
	# Extra safety: Ensure 'item' is a dictionary before reading keys
	if isinstance(item, dict) and "label" in item and "score" in item:
	clean_label = LABEL_MAPPING.get(item["label"], item["label"])
	formatted_results[clean_label] = float(item["score"])

	# Fallback if parsing completely fails to find labels
	if not formatted_results:
	return {"Unknown Error": 1.0}

	return formatted_results

	# Define Gradio Interface
	demo = gr.Interface(
	fn=detect_offensive_speech,
	inputs=gr.Textbox(placeholder="Enter text here...", lines=3),
	outputs=gr.Label(label="Classification Results", num_top_classes=2),
	title="Offensive Speech Detection",
	description="A local, optimized text-classification pipeline running on Hugging Face weights."
	)

	if __name__ == "__main__":
	demo.launch()