wilsonchang17/scamshield-api
0
1import os2import re3import json4import torch5import gradio as gr6from peft import PeftModel7from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria, StoppingCriteriaList8from dotenv import load_dotenv9 10# MODIFIED: Added StoppingCriteria to the imports.11 12# MODIFIED: Defined the StoppingCriteria class.13# This is the correct way to tell the model when to stop generating.14# It checks on-the-fly if the generated text ends with a stop string.15class StopOnJSON(StoppingCriteria):16 """17 Custom StoppingCriteria to stop generation when a complete JSON object is formed.18 """19 def __init__(self, stop_strings: list[str], tokenizer: AutoTokenizer):20 self.tokenizer = tokenizer21 self.stop_strings = stop_strings22 23 def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:24 # This function is called after each new token is generated.25 # It decodes the entire sequence to check for a stop condition.26 decoded_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)27 for stop_string in self.stop_strings:28 if decoded_text.endswith(stop_string):29 return True # Signal to stop generation30 return False # Signal to continue generation31 32class ScamDetector:33 """34 Detect scam messages and explain why, using a LoRA-adapted Llama-8B model.35 """36 def __init__(37 self,38 model_id: str = "meta-llama/Meta-Llama-3.1-8B-Instruct",39 adapter_path: str = "models/checkpoint-llama8",40 hf_token_env: str = "HF_TOKEN"41 ):42 self.model_id = model_id43 self.adapter_path = adapter_path44 self.hf_token_env = hf_token_env45 self.tokenizer = None46 self.model = None47 self.main_device = "cpu"48 49 # -------------------------------------------------------------------------50 def load_model(self):51 """52 Load tokenizer, base model, and LoRA adapter.53 Uses device_map='auto' so layers are automatically sharded across GPU/CPU.54 """55 load_dotenv()56 hf_token = os.getenv(self.hf_token_env)57 58 # 1. tokenizer59 self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, use_auth_token=hf_token)60 61 # 2. base model (auto device map ↔ accelerate hooks)62 base_model = AutoModelForCausalLM.from_pretrained(63 self.model_id,64 torch_dtype=torch.float16,65 device_map="auto",66 low_cpu_mem_usage=True,67 use_auth_token=hf_token,68 )69 70 # 3. attach LoRA adapter71 self.model = PeftModel.from_pretrained(72 base_model,73 self.adapter_path,74 torch_dtype=torch.float16,75 use_auth_token=hf_token,76 )77 self.model.eval()78 79 # 4. record "main" device for input tensors80 self.main_device = next(self.model.parameters()).device81 print(f"✅ Model loaded. Main device: {self.main_device}")82 83 # -------------------------------------------------------------------------84 def predict(self, message: str) -> tuple[str, str]:85 """Return ('Yes'|'No'|'Error', explanation_or_error)."""86 if self.model is None:87 raise RuntimeError("Model not loaded. Call load_model() first.")88 89 prompt = (90 "You are a scam-detection assistant.\n"91 "Respond ONLY with 'yes' (scam) or 'no' (not scam).\n\n"92 f"Message: {message}\n"93 "Answer:"94 )95 try:96 inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)97 inputs = {k: v.to(self.main_device) for k, v in inputs.items()}98 99 outputs = self.model.generate(100 **inputs,101 max_new_tokens=10,102 temperature=0.0,103 do_sample=False,104 top_k=1105 )106 full_answer = self.tokenizer.decode(107 outputs[0][inputs["input_ids"].size(1):], skip_special_tokens=True108 ).strip().lower()109 110 matches = re.findall(r'\b(yes|no|scam)\b', full_answer)111 if matches:112 final_decision = matches[-1].lower()113 if (final_decision == 'yes' or final_decision == 'scam'):114 return final_decision, self.explain(message)115 return final_decision, "Not a scam"116 return "No", "Not a scam"117 except Exception as e:118 return "Error", str(e)119 120 # -------------------------------------------------------------------------121 def explain(self, message: str) -> str:122 """Return bullet-point explanation for confirmed scam message."""123 if self.model is None:124 raise RuntimeError("Model not loaded. Call load_model() first.")125 126 prompt = (127 "You are a text-processing API endpoint. Your ONLY function is to return a single, valid JSON object based on the provided text. Do not output any other text, explanations, or markdown.\n\n"128 "<INSTRUCTIONS>\n"129 "1. Analyze the text inside the <SCAM_MESSAGE> tag.\n"130 "2. Identify exactly three distinct reasons why the message is a scam.\n"131 "3. Generate a JSON object that strictly adheres to the schema defined in the <JSON_SCHEMA> tag.\n"132 "4. Each reason must be a string of 50 words or less.\n"133 "</INSTRUCTIONS>\n\n"134 "<JSON_SCHEMA>\n"135 "{\n"136 ' "explanation": [\n'137 ' "string",\n'138 ' "string",\n'139 ' "string"\n'140 " ]\n"141 "}\n"142 "</JSON_SCHEMA>\n\n"143 "<SCAM_MESSAGE>\n"144 f"{message}\n"145 "</SCAM_MESSAGE>\n\n"146 "JSON_OUTPUT:\n"147 )148 inputs = self.tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)149 inputs = {k: v.to(self.main_device) for k, v in inputs.items()}150 151 # MODIFIED: Instantiate our custom stopping criteria.152 stop_criteria = StopOnJSON(stop_strings=["}", "}\n"], tokenizer=self.tokenizer)153 154 # MODIFIED: The generate call now uses the correct `stopping_criteria` parameter.155 # The invalid `stop` parameter has been removed.156 out = self.model.generate(157 **inputs,158 max_new_tokens=200,159 do_sample=False,160 temperature=0.1,161 pad_token_id=self.tokenizer.eos_token_id,162 stopping_criteria=StoppingCriteriaList([stop_criteria])163 )164 165 raw = self.tokenizer.decode(out[0][inputs["input_ids"].size(1):], skip_special_tokens=True).strip()166 print("Raw output from model:", raw) # It's good practice to log the raw output.167 168 # MODIFIED: Refined the post-processing to be a robust safety net.169 # It now looks for the first '{' and the first '}' that follows it.170 try:171 start = raw.find("{")172 end = raw.find("}", start) + 1173 if start != -1 and end != 0: # Ensure both braces were found174 raw = raw[start:end]175 else:176 # If no valid JSON structure is found, return the raw output as-is.177 return raw178 179 data = json.loads(raw)180 items = data.get("explanation", [])181 return "\n".join(f"• {item}" for item in items if item) # Added a check for empty items182 except (json.JSONDecodeError, TypeError, ValueError):183 return raw # If JSON parsing fails, return the cleaned-up raw string.184 185# =============================================================================186# Gradio interface187# =============================================================================188detector = ScamDetector()189print("⏳ Loading model...")190try:191 detector.load_model()192except Exception as e:193 print(f"❌ Model loading failed: {e}")194 detector = None195 196# ---------- API wrapper functions ----------197def predict_only(message: str):198 if detector is None or detector.model is None:199 return "Error", "Model not loaded"200 return detector.predict(message)201 202def explain_only(message: str):203 if detector is None or detector.model is None:204 return "Error: Model not loaded"205 return detector.explain(message)206 207# ---------- Gradio Blocks UI ----------208with gr.Blocks(title="ScamShield API") as demo:209 gr.Markdown("### ScamShield — Predict & Explain")210 211 inp = gr.Textbox(lines=4, placeholder="Paste a message here...")212 lbl = gr.Textbox(label="Label (Yes / No)")213 exp = gr.Textbox(label="Explanation")214 215 gr.Button("Predict (Yes/No + reason)").click(216 predict_only, inputs=inp, outputs=[lbl, exp], api_name="predict"217 )218 gr.Button("Explain only (already scam)").click(219 explain_only, inputs=inp, outputs=exp, api_name="explain"220 )221 222if __name__ == "__main__":223 demo.launch(share=True)224