wilsonchang17/scamshield-api

1import os2import re3import time4import torch5import gradio as gr6from peft import PeftModel7from transformers import AutoTokenizer, AutoModelForCausalLM8from dotenv import load_dotenv9 10load_dotenv()11 12MODEL_ID = "meta-llama/Meta-Llama-3.1-8B-Instruct"13ADAPTER_MODEL_PATH = "models/checkpoint-llama8"14HF_TOKEN = os.getenv("HF_TOKEN")15 16print("Loading tokenizer...")17tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)18print("Tokenizer loaded.")19 20print("Loading base model on GPU (auto device map)...")21base_model = AutoModelForCausalLM.from_pretrained(22    MODEL_ID,23    torch_dtype=torch.float16,24    device_map="auto",  # Automatically place layers on GPU(s)25    use_auth_token=HF_TOKEN26)27print("Base model loaded.")28 29print("Loading adapter weights...")30model = PeftModel.from_pretrained(base_model, ADAPTER_MODEL_PATH, use_auth_token=HF_TOKEN)31print("Adapter loaded.")32 33# Optional warm-up procedure to ensure the model is fully ready34print("Warming up the model...")35dummy_prompt = "Hello, this is a warm-up prompt."36inputs = tokenizer(dummy_prompt, return_tensors="pt", truncation=True, max_length=512)37# Move inputs to the same device as the model38inputs = {key: value.to(model.device) for key, value in inputs.items()}39with torch.no_grad():40    _ = model.generate(41        **inputs,42        max_new_tokens=5,43        do_sample=True,44        temperature=0.7,45        top_p=0.9,46        top_k=50,47        pad_token_id=tokenizer.eos_token_id48    )49print("Model warm-up completed.")50 51def classify_message_llm(message):52    print("Received message:", message)53    prompt = (54    "You are a scam‑detection assistant.\n"55    "Respond ONLY with 'yes' (scam) or 'no' (not scam).\n\n"56    f"Message: {message}\n"57    "Answer:"58    )59    60    try:61        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)62        # Move inputs to the same device as the model63        inputs = {key: value.to(model.device) for key, value in inputs.items()}64        print("Generating...", flush=True)65        model.eval()66        with torch.no_grad():67            outputs = model.generate(68                **inputs,69                max_new_tokens=10,70                temperature=0.0,71                do_sample=False,72                top_k=173            )74        print("Generation completed", flush=True)75        generated_tokens = outputs[0][inputs['input_ids'].size(1):]76        full_answer = tokenizer.decode(generated_tokens, skip_special_tokens=True).strip().lower()77        print("Full answer:", full_answer)78        matches = re.findall(r'\b(yes|no|scam)\b', full_answer)79        if matches:80            final_decision = matches[-1]81            return 1 if (final_decision == 'yes' or final_decision == 'scam') else 082        else:83            return -184    except Exception as e:85        print(f"LLM evaluation error: {e}")86        return -187 88demo = gr.Interface(89    fn=classify_message_llm,90    inputs=gr.Textbox(lines=5, placeholder="Enter a message to check if it's a scam"),91    outputs="text",92    title="Scam Detection with Fine-tuned LLaMA 3.1 8B"93)94 95if __name__ == "__main__":96    demo.launch(share=True)97