from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

def predict_code(code_snippet):
    # 1. Load the "Brain" we saved
    model_path = "saved_model"
    tokenizer = AutoTokenizer.from_pretrained(model_path)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)

    # 2. Convert the code snippet into numbers (Tokens)
    inputs = tokenizer(code_snippet, return_tensors="pt", truncation=True, max_length=128)

    # 3. Run it through the AI without calculating gradients (faster)
    with torch.no_grad():
        outputs = model(**inputs)
        # Convert raw numbers (logits) into percentages (probabilities)
        probs = F.softmax(outputs.logits, dim=-1)

    # 4. Show the result
    safe_prob = probs[0][0].item()
    vuln_prob = probs[0][1].item()

    print("-" * 50)
    print(f"CODE BEING TESTED:\n{code_snippet}")
    print("-" * 50)
    print(f"🛡️  Safe Probability: {safe_prob:.2%}")
    print(f"⚠️  Vulnerable Probability: {vuln_prob:.2%}")
    
    if vuln_prob > 0.5:
        print("\nRESULT: 🚨 VULNERABLE CODE DETECTED!")
    else:
        print("\nRESULT: ✅ CODE LOOKS SAFE.")
    print("-" * 50)

if __name__ == "__main__":
    # Test with a dangerous example (SQL Injection)
    my_code = 'query = f"SELECT * FROM users WHERE id = {user_input}"'
    predict_code(my_code)