Spaces:
Sleeping
Sleeping
File size: 4,662 Bytes
f6eebd3 0bbe315 f6eebd3 e1ab938 f6eebd3 bc4efd2 02cb142 f6eebd3 3d0613b 56cde6f 3d0613b f6eebd3 | 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 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
# ---------------------------------------------------------------------------
# CONFIGURATION
# ---------------------------------------------------------------------------
MODEL_PATH = "./model"
# Automatically detect hardware
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Running on: {device}")
# ---------------------------------------------------------------------------
# MODEL LOADING
# ---------------------------------------------------------------------------
try:
tokenizer = AutoTokenizer.from_pretrained(
MODEL_PATH,
use_fast=False, # Essential for SentencePiece tokenizers (like Gemma/Llama)
trust_remote_code=True
)
# Load model with appropriate precision
# GPU = float16 (faster, less VRAM)
# CPU = float32 (required for compatibility on basic CPU Spaces)
torch_dtype = torch.float16 if device == "cuda" else torch.float32
model = AutoModelForCausalLM.from_pretrained(
MODEL_PATH,
torch_dtype=torch_dtype,
trust_remote_code=True
).to(device)
model.eval()
print("Model loaded successfully.")
except Exception as e:
print(f"FATAL ERROR loading model: {e}")
raise e
# ---------------------------------------------------------------------------
# INFERENCE FUNCTION
# ---------------------------------------------------------------------------
def classify_spam(text):
# 1. Input Validation
if not text or not text.strip():
return "⚠️ Please enter a message."
# 2. Prepare Prompt
messages = [
{"role": "user", "content": text},
]
try:
# Apply template (returns a string)
prompt_str = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)
# Convert string to PyTorch tensors and move to device
model_inputs = tokenizer(prompt_str, return_tensors="pt").to(device)
# 3. Generate
with torch.no_grad():
outputs = model.generate(
**model_inputs,
max_new_tokens=20, # Keep strictly short for classification
do_sample=False, # STRICTLY REQUIRED: Deterministic output (no randomness)
pad_token_id=tokenizer.eos_token_id
)
# 4. Decode and Clean
# The model returns [Prompt + Answer]. We slice off the prompt.
input_length = model_inputs["input_ids"].shape[1]
generated_tokens = outputs[0][input_length:]
result = tokenizer.decode(generated_tokens, skip_special_tokens=True)
# FIX: Manually remove the specific tag and whitespace
cleaned_result = result.replace("<end_of_turn>", "").strip()
return cleaned_result
except Exception as e:
return f"Error during inference: {str(e)}"
# ---------------------------------------------------------------------------
# GRADIO INTERFACE
# ---------------------------------------------------------------------------
with gr.Blocks(title="Selling Intent Classifier") as demo:
gr.Markdown("# 📧 Selling Intent Classifier")
gr.Markdown("Enter a youtube comment section message selling or asking to subscribe channel")
with gr.Row():
inp = gr.Textbox(label="Message Content", placeholder="Paste text here...", lines=4)
out = gr.Textbox(label="Classification Result")
btn = gr.Button("Classify", variant="primary")
# Connecting the button to the function
btn.click(fn=classify_spam, inputs=inp, outputs=out)
# --- ADDED EXAMPLES SECTION HERE ---
gr.Examples(
examples=[
["ye teacher bakar padhata hai mujha sub karlo pwalternative on sc"],
["""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 ............
Channel name : - Edu Notes
(See logo of ytube channel )"""],
["ye sub toh bahut hard hai maths is tough"],
["Click here for a free prize and subscribe to my channel!"]
],
inputs=inp,
label="Try these examples:"
)
# ---------------------------------------------------------------------------
# LAUNCH
# ---------------------------------------------------------------------------
if __name__ == "__main__":
demo.queue().launch() |