Create main.py

main.py

@@ -0,0 +1,69 @@
+from flask import Flask, request, jsonify
+import torch
+from transformers import RobertaTokenizer
+import os
+from transformers import RobertaForSequenceClassification
+import torch.serialization
+import torch
+from transformers import RobertaTokenizer, RobertaForSequenceClassification, Trainer, TrainingArguments
+from torch.utils.data import Dataset
+import pandas as pd
+from sklearn.model_selection import train_test_split
+import numpy as np
+# Initialize Flask app
+app = Flask(__name__)
+
+# Load the trained model and tokenizer
+tokenizer = RobertaTokenizer.from_pretrained("microsoft/codebert-base")
+torch.serialization.add_safe_globals([RobertaForSequenceClassification])
+
+model = torch.load("model.pth", map_location=torch.device('cpu'), weights_only=False)  # Load the trained model
+
+# Ensure the model is in evaluation mode
+model.eval()
+
+
+@app.route("/")
+def home():
+    return request.url
+
+
+# @app.route("/predict", methods=["POST"])
+@app.route("/predict")
+def predict():
+    print("Received code:", request.get_json()["code"])
+    code = request.get_json()["code"]
+    # Load saved weights and config
+    checkpoint = torch.load("codebert_vulnerability_scorer.pth")
+    config = RobertaConfig.from_dict(checkpoint['config'])
+
+    # Rebuild the model with correct architecture
+    model = RobertaForSequenceClassification(config)
+    model.load_state_dict(checkpoint['model_state_dict'])
+    model.eval()
+
+    # Load tokenizer
+    tokenizer = RobertaTokenizer.from_pretrained('./tokenizer')
+
+    # Prepare input
+    inputs = tokenizer(
+        code,
+        truncation=True,
+        padding='max_length',
+        max_length=512,
+        return_tensors='pt'
+    )
+
+    # Make prediction
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    score = torch.sigmoid(outputs.logits).item()
+    return score
+    
+
+
+    
+# Run the Flask app
+if __name__ == "__main__":
+     app.run(host="0.0.0.0", port=7860)