Building a vulnerable code detector with TorchSharp

This is the fifth article in the series where we are teaching a machine to find security flaws in code. The trip so far: machines learn from examples (article 1), transformers read understanding context (article 2), a vulnerability is a pattern born from too much trust in user input (article 3), and CodeBERT is the code reader we chose for the job (article 4).

The time has come to assemble the detector. And I am going to do it the boldest way: everything in C#. The training, the execution, everything.

The assembly has four pieces: the vocabulary, the model, the specialization and the verdict. Let's get to them.

Piece 1: the vocabulary (tokenizer)

No model reads letters. Every transformer sees the world as numbers, and the piece that converts text into numbers is called the tokenizer. It breaks the text into little pieces (tokens) and swaps each piece for its number in a giant table, the vocabulary. CodeBERT has its own vocabulary, built during its "college" years, and we need to use exactly the same one, otherwise it's like talking to it in the wrong language.

In .NET, the Microsoft.ML.Tokenizers package solves this. The vocabulary files are published along with the model, on CodeBERT's official page on Hugging Face, and there are three of them: vocab.json is the table linking each token to its number, merges.txt holds the rules for how to break the text into little pieces, and dict.txt is an auxiliary frequency map this tokenizer asks for. Once you download the three, just load them. And as an example, let's convert exactly the kind of snippet the detector will analyze, a SQL query built in C#:

// Loads CodeBERT's official vocabulary (the 3 files downloaded from Hugging Face)
var tokenizer = new EnglishRoberta("vocab.json", "merges.txt", "dict.txt");

// The exact kind of snippet the detector will analyze
var suspiciousCode = """var sql = "SELECT * FROM Users WHERE Name = '" + name + "'";""";

// Converts the code into numbers, the only language the model understands
var tokenIds = tokenizer.EncodeToIds(suspiciousCode);
// tokenIds: [318, 1640, 22872, 4, ...]

Notice the three quotes around the snippet: those are modern C#'s raw strings, which let us keep code full of quotes inside a string without escaping anything. They will show up again later.

To you, these numbers mean nothing. To CodeBERT, they are the whole sentence.

Piece 2: the model (and the borrowed knowledge)

Now the central piece. TorchSharp is the .NET version of PyTorch, the most used deep learning library in the world, and it is with it that we assemble the model's structure. That structure has two parts, and it's worth separating them clearly in your head.

The first is the body of CodeBERT: the entry that receives the tokenizer's numbers and the stack of attention layers (the ones from article 2, where everyone looks at everyone). This is where all the ability to "understand" code lives.

The second is the classification layer: a new, small layer we place on top of the body. Its job is to take everything the body understood about the code snippet and summarize it into two outputs: the "vulnerable" score and the "safe" score. It is born blank, knowing nothing, and it is mostly this layer that our training will teach.

// The structure: CodeBERT's pre-trained body + the classification layer (2 outputs)
var model = new CodeBertClassifier(labelCount: 2);

// Dresses the body with the pre-trained knowledge published by Microsoft
// (the weights file is read by the TorchSharp.PyBridge package)
model.load_py("codebert-base.bin");

That second line deserves a paragraph. The codebert-base.bin file contains CodeBERT's weights: millions of little numbers adjusted during pre-training which are, literally, the model's knowledge. It is the result of the "six years of college" we are not going to pay for. When TorchSharp.PyBridge reads that file and fills our body with those values, the entire CodeBERT comes to exist, alive, inside a .NET process. Only the classification layer on top stays blank, waiting for our lesson.

And here I make a scriptwriter's confession: in this article, the attention layers are closed boxes. We plug them in, stack them and use them, without looking inside. It's not laziness, it's suspense: opening those boxes and assembling every gear of them from scratch is exactly the subject of the last article of the series. Here the focus is getting the detector to work.

Piece 3: the specialization (the training, in C#)

The moment the whole series has been preparing: teaching. Remember article 1: labeled examples. We build a collection of pairs in C#, always with both versions of the same snippet.

First, the wrong-way example, which we label as vulnerable. It is the exact pattern from article 3, the user's data glued into the middle of the command:

List<LabeledExample> trainingExamples = [];

// The pattern from article 3: the user's data glued into the middle of the command
trainingExamples.Add(new LabeledExample(
    Code: """var sql = "SELECT * FROM Users WHERE Name = '" + name + "'";""",
    Label: Label.Vulnerable));

And the corrected version of the same snippet, labeled as safe, where the data travels separately from the command, as a parameter:

// The same query done right: the data travels separately, as a parameter
trainingExamples.Add(new LabeledExample(
    Code: """var sql = "SELECT * FROM Users WHERE Name = @name";""",
    Label: Label.Safe));

We repeat this in dozens of variations, covering the classics: injection, file paths built with user input, and so on. Half vulnerable, half safe, and a separate part the model never sees during training (the final exam with brand-new questions, remember?).

And the training itself? It is a three-step ritual that repeats, and I find it enormously elegant:

const int epochs = 5;              // how many times the model rereads the whole material
const double learningRate = 2e-5;  // the size of each adjustment (small on purpose)

var optimizer = torch.optim.AdamW(model.parameters(), lr: learningRate);
var lossFunction = CrossEntropyLoss();

for (int epoch = 1; epoch <= epochs; epoch++)
{
    foreach (var example in trainingExamples)
    {
        var prediction = model.forward(tokenizer.ToTensor(example.Code));  // the guess
        var loss = lossFunction.forward(prediction, example.LabelTensor);  // the error grade

        optimizer.zero_grad();  // clears the notes from the previous round
        loss.backward();        // finds out how much each piece contributed to the error
        optimizer.step();       // adjusts each piece a little bit in the right direction
    }
}

Translating the ritual: the model makes a guess, the loss function gives the grade (how far the guess was from the right answer), and the optimizer makes the correction, slightly adjusting each parameter in the direction that would have reduced the error. Guess, grade, correction. Thousands of times. It is literally studying by trial and error, and that is how all deep learning on the planet works, from our detector to the giant models.

The five epochs are the five rereads of the material. And the cost? With our education-sized collection, a machine with an ordinary gaming graphics card solves it in minutes. Without one, on the processor alone, it's worth a longer coffee break. No cloud, no bill to pay.

Piece 4: the verdict

Once trained, the model becomes what we always wanted: a function that reads code and answers. Execution is the same path as training, just without the correction: convert the code into numbers, pass it through the model, and the output is two probabilities.

var suspiciousCode = """var sql = "DELETE FROM Logs WHERE Id = " + id;""";

var scores = model.forward(tokenizer.ToTensor(suspiciousCode));
var probabilities = Softmax(scores);   // the scores become percentages that add up to 100%

Console.WriteLine($"Probability of vulnerable: {probabilities.Vulnerable:P0}");
// Probability of vulnerable: 93%

And that's it. A transformer specialized in security, trained by you, running locally in a .NET console app, answering in milliseconds. Not a single line of your code left your machine.

The dose of reality

To keep the house's standard of honesty, two quick reminders. CodeBERT did not study C# in college (its six languages did not include our idiom), and that is why that collection of C# examples we built in Piece 3 matters so much: it is what introduces C# to it. And a detector trained with an educational collection serves to demonstrate the concept and recognize the classic patterns, not to replace the defenses from article 3. It is the fourth layer, the one that reads with context and costs almost nothing per analysis, adding to the others, never in their place.

Conclusion

What we assembled:

Tokenizer: CodeBERT's vocabulary loaded in C# with Microsoft.ML.Tokenizers
Model: the structure assembled in TorchSharp, dressed with the pre-trained weights via TorchSharp.PyBridge
Training: the guess, grade and correction ritual, running on your machine, in C#, for free
Verdict: code goes in, probability of vulnerability comes out, in milliseconds

The detector is standing. But a deliberate debt was left along the way: those closed boxes, the attention layers we stacked without opening. In the next article, the last of the series, we pay that debt with interest: we are going to build a transformer from scratch, piece by piece, in C#, and you will see with your own eyes the gear that moves everything this series has told. It is the article I most want to write.

Blog

Series