| import gradio as gr |
| import json |
| import time |
| import spaces |
| from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline |
|
|
| MODEL_PATH = 'berkeruveyik/toxic-speech-finetune-with-gemma-3-1b-v1' |
|
|
| |
| loaded_model = AutoModelForCausalLM.from_pretrained( |
| MODEL_PATH, |
| torch_dtype='auto', |
| device_map='auto', |
| attn_implementation='eager' |
| ) |
|
|
| tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) |
|
|
| loaded_model_pipeline = pipeline( |
| 'text-generation', |
| model=loaded_model, |
| tokenizer=tokenizer |
| ) |
|
|
| @spaces.GPU |
| def pred_on_text(input_text): |
| """Generate prediction from input text""" |
| start_time = time.time() |
|
|
| raw_output = loaded_model_pipeline( |
| text_inputs=[{'role': 'user', 'content': input_text}], |
| max_new_tokens=256 |
| ) |
|
|
| end_time = time.time() |
| total_time = round(end_time - start_time, 4) |
|
|
| generated_text = raw_output[0]['generated_text'][1]['content'] |
|
|
| return generated_text, raw_output, total_time |
|
|
| def parse_generated_text(text): |
| """Parse the generated text and format it nicely""" |
| try: |
| data = json.loads(text) |
| return data |
| except: |
| try: |
| text = text.strip() |
| if text.startswith('{') and text.endswith('}'): |
| data = eval(text) |
| return data |
| except: |
| pass |
| return {"raw_output": text} |
|
|
| def format_output(input_text, parsed_output, total_time): |
| """Format output as readable text with each field on new line""" |
| output_lines = [] |
| output_lines.append(f"📝 Input: {input_text}") |
| output_lines.append("") |
| output_lines.append("━" * 50) |
| output_lines.append("") |
|
|
| if "is_toxic" in parsed_output: |
| emoji = "🚨" if parsed_output["is_toxic"] else "✅" |
| output_lines.append(f"{emoji} is_toxic: {parsed_output['is_toxic']}") |
|
|
| if "label" in parsed_output: |
| output_lines.append(f"🏷️ label: {parsed_output['label']}") |
|
|
| if "tags" in parsed_output: |
| output_lines.append(f"🔖 tags: {parsed_output['tags']}") |
|
|
| if "reason" in parsed_output: |
| output_lines.append(f"💬 reason: {parsed_output['reason']}") |
|
|
| if "severity" in parsed_output: |
| output_lines.append(f"⚡ severity: {parsed_output['severity']}") |
|
|
| if "raw_output" in parsed_output: |
| output_lines.append(f"📄 raw_output: {parsed_output['raw_output']}") |
|
|
| output_lines.append("") |
| output_lines.append("━" * 50) |
| output_lines.append(f"⏱️ processing_time: {total_time} seconds") |
|
|
| return "\n".join(output_lines) |
|
|
| def gradio_predict(input_text): |
| """Wrapper function for Gradio""" |
| if not input_text.strip(): |
| return "Please enter some text." |
|
|
| generated_text, raw_output, total_time = pred_on_text(input_text) |
| parsed_output = parse_generated_text(generated_text) |
|
|
| formatted_output = format_output(input_text, parsed_output, total_time) |
|
|
| return formatted_output |
|
|
| |
| demo = gr.Interface( |
| fn=gradio_predict, |
| inputs=gr.Textbox( |
| label="Input Text", |
| placeholder="Enter your text here...", |
| lines=3 |
| ), |
| outputs=gr.Textbox( |
| label="Model Output", |
| lines=12 |
| ), |
| title="🤖 Toxic Speech Classifier", |
| description="Analyze whether a given text contains toxic, insulting, or harmful language using a fine-tuned Gemma3 model.", |
| examples=[ |
| ["You are absolutely worthless and no one will ever love you."], |
| ["I hope you get hit by a bus, you disgusting excuse for a person."], |
| ["The weather today is really nice, I enjoyed my walk in the park."], |
| ["Shut up you brainless moron, nobody asked for your stupid opinion."], |
| ["Thank you for your help, I really appreciate everything you did."], |
| ["You are such a pathetic loser, get out of my sight."], |
| ["I just finished reading a great book, it was very inspiring."], |
| ], |
| theme=gr.themes.Soft() |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
