from flask import Flask, request, jsonify
import torch
from transformers import RobertaTokenizer, RobertaForSequenceClassification, RobertaConfig
import os

app = Flask(__name__)

# Load model and tokenizer
def load_model():
    # Load saved config and weights
    # Updated to 'codebert_vulnerability_scorer.pth' as per user's request
    checkpoint = torch.load("codebert_vulnerability_scorer.pth", map_location=torch.device('cpu'))
    config = RobertaConfig.from_dict(checkpoint['config'])

    # Initialize model with loaded config
    model = RobertaForSequenceClassification(config)
    model.load_state_dict(checkpoint['model_state_dict'])
    model.eval()
    return model

# Load components
try:
    # Updated to './tokenizer_vulnerability' as per user's request
    tokenizer = RobertaTokenizer.from_pretrained("./tokenizer_vulnerability")
    model = load_model()
    print("Model and tokenizer loaded successfully!")
except Exception as e:
    print(f"Error loading model: {str(e)}")

@app.route("/")
def home():
    return request.url

@app.route("/predict", methods=["POST"])  
def predict():
    try:
        data = request.get_json()
        if not data or "code" not in data:
            return jsonify({"error": "Missing 'code' in request body"}), 400

        code = data["code"]

        # Tokenize input
        inputs = tokenizer(
            code,
            truncation=True,
            padding='max_length',
            max_length=512,
            return_tensors='pt'
        )

        # Make prediction
        with torch.no_grad():
            outputs = model(**inputs)

        # Apply sigmoid and format score
        score = torch.sigmoid(outputs.logits).item()

        return jsonify({
            "score": round(score, 4)
        })

    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=7860)