main.py

from flask import Flask, request, jsonify
import torch
from transformers import RobertaTokenizer, RobertaForSequenceClassification
import os
import gc
from functools import lru_cache

app = Flask(__name__)

model = None
tokenizer = None
device = None

def setup_device():
    if torch.cuda.is_available():
        return torch.device('cuda')
    elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
        return torch.device('mps')
    else:
        return torch.device('cpu')

def load_tokenizer():
    try:
        tokenizer = RobertaTokenizer.from_pretrained('./tokenizer_vulnerability')
        tokenizer.model_max_length = 512
        return tokenizer
    except Exception as e:
        print(f"Error loading tokenizer: {e}")
        try:
            return RobertaTokenizer.from_pretrained('microsoft/codebert-base')
        except Exception as e2:
            print(f"Fallback tokenizer failed: {e2}")
            return None

def load_model():
    global device
    device = setup_device()
    print(f"Using device: {device}")
    
    try:
        checkpoint = torch.load("codebert_vulnerability_scorer.pth", map_location=device)
        
        if 'config' in checkpoint:
            from transformers import RobertaConfig
            config = RobertaConfig.from_dict(checkpoint['config'])
            model = RobertaForSequenceClassification(config)
        else:
            model = RobertaForSequenceClassification.from_pretrained(
                'microsoft/codebert-base',
                num_labels=1
            )
        
        if 'model_state_dict' in checkpoint:
            model.load_state_dict(checkpoint['model_state_dict'])
        else:
            model.load_state_dict(checkpoint)
            
        model.to(device)
        model.eval()
        
        if device.type == 'cuda':
            model.half()
        
        return model
        
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

def cleanup_gpu_memory():
    if device and device.type == 'cuda':
        torch.cuda.empty_cache()
    gc.collect()

try:
    print("Loading tokenizer...")
    tokenizer = load_tokenizer()
    if tokenizer:
        print("Tokenizer loaded successfully!")
    else:
        print("Failed to load tokenizer!")
    
    print("Loading model...")
    model = load_model()
    if model:
        print("Model loaded successfully!")
    else:
        print("Failed to load model!")
    
except Exception as e:
    print(f"Error during initialization: {str(e)}")
    tokenizer = None
    model = None

@app.route("/", methods=['GET'])
def home():
    return jsonify({
        "message": "CodeBERT Vulnerability Evalutor API",
        "status": "Model loaded" if model is not None else "Model not loaded",
        "device": str(device) if device else "unknown",
        "endpoints": {
            "/predict": "POST with JSON body containing 'codes' array"
        }
    })

@app.route("/predict", methods=['POST'])
def predict_batch():
    try:
        if model is None or tokenizer is None:
            return jsonify({"error": "Model not loaded properly"}), 500
            
        data = request.get_json()
        if not data or 'codes' not in data:
            return jsonify({"error": "Missing 'codes' field in JSON body"}), 400
            
        codes = data['codes']
        if not isinstance(codes, list) or len(codes) == 0:
            return jsonify({"error": "'codes' must be a non-empty array"}), 400
        
        if len(codes) > 100:
            return jsonify({"error": "Too many codes. Maximum 100 allowed."}), 400
        
        validated_codes = []
        for i, code in enumerate(codes):
            if not isinstance(code, str):
                return jsonify({"error": f"Code at index {i} must be a string"}), 400
            if len(code.strip()) == 0:
                validated_codes.append("# empty code")
            elif len(code) > 50000:
                return jsonify({"error": f"Code at index {i} too long. Maximum 50000 characters."}), 400
            else:
                validated_codes.append(code.strip())
        
        if len(validated_codes) == 1:
            score = predict_vulnerability_with_chunking(validated_codes[0])
            cleanup_gpu_memory()
            return jsonify({"results": [{"score": 1.0 - score}]})
        
        batch_size = min(len(validated_codes), 16)
        results = []
        
        try:
            for i in range(0, len(validated_codes), batch_size):
                batch = validated_codes[i:i+batch_size]
                
                long_codes = []
                short_codes = []
                long_indices = []
                short_indices = []
                
                for idx, code in enumerate(batch):
                    try:
                        tokens = tokenizer.encode(code, add_special_tokens=False, max_length=1000, truncation=True)
                        if len(tokens) > 450:
                            long_codes.append(code)
                            long_indices.append(i + idx)
                        else:
                            short_codes.append(code)
                            short_indices.append(i + idx)
                    except Exception as e:
                        print(f"Tokenization error for code {i + idx}: {e}")
                        short_codes.append(code)
                        short_indices.append(i + idx)
                
                batch_scores = [0.0] * len(batch)
                
                if short_codes:
                    try:
                        short_scores = predict_vulnerability_batch(short_codes)
                        for j, score in enumerate(short_scores):
                            local_idx = short_indices[j] - i
                            batch_scores[local_idx] = score
                    except Exception as e:
                        print(f"Batch prediction error: {e}")
                        for j in range(len(short_codes)):
                            local_idx = short_indices[j] - i
                            batch_scores[local_idx] = 0.0
                
                for j, code in enumerate(long_codes):
                    try:
                        score = predict_vulnerability_with_chunking(code)
                        local_idx = long_indices[j] - i
                        batch_scores[local_idx] = score
                    except Exception as e:
                        print(f"Chunking prediction error: {e}")
                        local_idx = long_indices[j] - i
                        batch_scores[local_idx] = 0.0
                
                for score in batch_scores:
                    results.append({"score": round(1.0 - score,4)})
                
                cleanup_gpu_memory()
        
        except Exception as e:
            cleanup_gpu_memory()
            raise e
        
        return jsonify({"results": results})
        
    except Exception as e:
        cleanup_gpu_memory()
        return jsonify({"error": f"Batch prediction error: {str(e)}"}), 500

def predict_vulnerability_with_chunking(code):
    try:
        if not code or len(code.strip()) == 0:
            return 0.0
            
        tokens = tokenizer.encode(code, add_special_tokens=False, max_length=2000, truncation=True)
        
        if len(tokens) <= 450:
            return predict_vulnerability(code)
        
        chunk_size = 400
        overlap = 50
        max_score = 0.0
        
        for start in range(0, len(tokens), chunk_size - overlap):
            end = min(start + chunk_size, len(tokens))
            chunk_tokens = tokens[start:end]
            
            try:
                chunk_code = tokenizer.decode(chunk_tokens, skip_special_tokens=True)
                if chunk_code.strip():
                    score = predict_vulnerability(chunk_code)
                    max_score = max(max_score, score)
            except Exception as e:
                print(f"Chunk processing error: {e}")
                continue
            
            if end >= len(tokens):
                break
        
        return max_score
        
    except Exception as e:
        print(f"Chunking error: {e}")
        return 0.0

def predict_vulnerability(code):
    try:
        if not code or len(code.strip()) == 0:
            return 0.0
            
        dynamic_length = min(max(len(code.split()) * 2, 128), 512)
        
        inputs = tokenizer(
            code,
            truncation=True,
            padding='max_length',
            max_length=dynamic_length,
            return_tensors='pt'
        )
        
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            if device.type == 'cuda':
                with torch.cuda.amp.autocast():
                    outputs = model(**inputs)
            else:
                outputs = model(**inputs)
        
        amplified_logits = 2.0 * outputs.logits
        score = torch.sigmoid(amplified_logits).cpu().item()
        return round(max(0.0, min(1.0, score)), 4)
        
    except Exception as e:
        print(f"Single prediction error: {e}")
        return 0.0

def predict_vulnerability_batch(codes):
    try:
        if not codes or len(codes) == 0:
            return []
            
        filtered_codes = [code if code and code.strip() else "# empty" for code in codes]
        
        max_len = max([len(code.split()) * 2 for code in filtered_codes if code])
        dynamic_length = min(max(max_len, 128), 512)
        
        inputs = tokenizer(
            filtered_codes,
            truncation=True,
            padding='max_length',
            max_length=dynamic_length,
            return_tensors='pt'
        )
        
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            if device.type == 'cuda':
                with torch.cuda.amp.autocast():
                    outputs = model(**inputs)
            else:
                outputs = model(**inputs)
        
        amplified_logits = 2.0 * outputs.logits
        scores = torch.sigmoid(amplified_logits).cpu().numpy()

        return [round(max(0.0, min(1.0, float(score))), 4) for score in scores.flatten()]
        
    except Exception as e:
        print(f"Batch prediction error: {e}")
        return [0.0] * len(codes)

@app.route("/health", methods=['GET'])
def health_check():
    return jsonify({
        "status": "healthy",
        "model_loaded": model is not None,
        "tokenizer_loaded": tokenizer is not None,
        "device": str(device) if device else "unknown"
    })

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