KushJaggi/spam_classifier

1import gradio as gr2import torch3from transformers import AutoTokenizer, AutoModelForCausalLM4 5# ---------------------------------------------------------------------------6# CONFIGURATION7# ---------------------------------------------------------------------------8 9MODEL_PATH = "./model" 10 11# Automatically detect hardware12device = "cuda" if torch.cuda.is_available() else "cpu"13print(f"Running on: {device}")14 15# ---------------------------------------------------------------------------16# MODEL LOADING17# ---------------------------------------------------------------------------18try:19    tokenizer = AutoTokenizer.from_pretrained(20        MODEL_PATH,21        use_fast=False, # Essential for SentencePiece tokenizers (like Gemma/Llama)22        trust_remote_code=True23    )24    25    # Load model with appropriate precision26    # GPU = float16 (faster, less VRAM)27    # CPU = float32 (required for compatibility on basic CPU Spaces)28    torch_dtype = torch.float16 if device == "cuda" else torch.float3229    30    model = AutoModelForCausalLM.from_pretrained(31        MODEL_PATH,32        torch_dtype=torch_dtype,33        trust_remote_code=True34    ).to(device)35    36    model.eval()37    print("Model loaded successfully.")38 39except Exception as e:40    print(f"FATAL ERROR loading model: {e}")41    raise e42 43# ---------------------------------------------------------------------------44# INFERENCE FUNCTION45# ---------------------------------------------------------------------------46def classify_spam(text):47    # 1. Input Validation48    if not text or not text.strip():49        return "⚠️ Please enter a message."50 51    # 2. Prepare Prompt52    messages = [53        {"role": "user", "content": text},54    ]55    56    try:57        # Apply template (returns a string)58        prompt_str = tokenizer.apply_chat_template(59            messages, 60            tokenize=False, 61            add_generation_prompt=True62        )63        64        # Convert string to PyTorch tensors and move to device65        model_inputs = tokenizer(prompt_str, return_tensors="pt").to(device)66        67        # 3. Generate68        with torch.no_grad():69            outputs = model.generate(70                **model_inputs,71                max_new_tokens=20,   # Keep strictly short for classification72                do_sample=False,     # STRICTLY REQUIRED: Deterministic output (no randomness)73                pad_token_id=tokenizer.eos_token_id74            )75 76        # 4. Decode and Clean77        # The model returns [Prompt + Answer]. We slice off the prompt.78        input_length = model_inputs["input_ids"].shape[1]79        generated_tokens = outputs[0][input_length:]80        81        result = tokenizer.decode(generated_tokens, skip_special_tokens=True)82        # FIX: Manually remove the specific tag and whitespace83        cleaned_result = result.replace("<end_of_turn>", "").strip()84        85        return cleaned_result86 87    except Exception as e:88        return f"Error during inference: {str(e)}"89 90# ---------------------------------------------------------------------------91# GRADIO INTERFACE92# ---------------------------------------------------------------------------93with gr.Blocks(title="Selling Intent Classifier") as demo:94    gr.Markdown("# 📧 Selling Intent Classifier")95    gr.Markdown("Enter a youtube comment section message selling or asking to subscribe channel")96    97    with gr.Row():98        inp = gr.Textbox(label="Message Content", placeholder="Paste text here...", lines=4)99        out = gr.Textbox(label="Classification Result")100    101    btn = gr.Button("Classify", variant="primary")102    103    # Connecting the button to the function104    btn.click(fn=classify_spam, inputs=inp, outputs=out)105 106    # --- ADDED EXAMPLES SECTION HERE ---107    gr.Examples(108        examples=[109            ["ye teacher bakar padhata hai mujha sub karlo pwalternative on sc"],110            ["""Hello Dosto / Bhaiyo / Bahno agar aapko Notes banaane me problem aa rahi hai ya aapko Notes Chaihiye school ke liye to ye ytube channel dekhiye .......  ya aapko Notes , Formula sheet , and exam related chize milengi ... jisse aapki Help ho sake  ............  111 112            Channel name : -  Edu Notes 113 114            (See logo of ytube channel )"""],115            ["ye sub toh bahut hard hai maths is tough"],116            ["Click here for a free prize and subscribe to my channel!"]117        ],118        inputs=inp,119        label="Try these examples:"120    )121 122# ---------------------------------------------------------------------------123# LAUNCH124# ---------------------------------------------------------------------------125if __name__ == "__main__":126    demo.queue().launch()