Maybe you even use the LLM to find vulnerable snippets at the beginning, but a multi class classifier or embedding model will be way better at runtime.
If data is the issue you can probably even generate vulnerabilities to create a synthetic dataset.
With enough data, you could train a classic ml model, or you could keep the llm in the inference pipeline, but is there another way?
1. Train an embedding model which forces “similar” inputs close together using triplet loss. Here “similar” can mean anything, but you would probably want to mark similar vulnerabilities as being similar.
2. If you have a fixed set of N vulnerabilities you can train a multi class classifier. Of course it’s a pain in the ass to add a new class later on.
3. For any particular vulnerability you could train a ranking model using hinge loss. This is what most industrial ranking and recommendation systems do.