import sys
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
from pathlib import Path

path = Path(sys.argv[1])
text = path.read_text(encoding="utf-8")

tokenizer = AutoTokenizer.from_pretrained("LeeSek/binary-dockerfile-model")
model = AutoModelForSequenceClassification.from_pretrained("LeeSek/binary-dockerfile-model")
model.eval()

inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=512)

with torch.no_grad():
    logits = model(**inputs).logits
    probs = torch.nn.functional.softmax(logits, dim=1).squeeze()

label = "GOOD" if torch.argmax(probs).item() == 0 else "BAD"
print(f"Prediction: {label} — Probabilities: good={probs[0]:.3f}, bad={probs[1]:.3f}")

📄 Step 2 — Create good and bad Dockerfile

Good:

docker

FROM node:18
WORKDIR /app
COPY . .
RUN npm install
CMD ["node", "index.js"]

Bad:

docker

FROM ubuntu:latest
RUN apt-get install python3
ADD . /app
WORKDIR /app
RUN pip install flask
CMD python3 app.py

▶️ Step 3 — Run the prediction

python test_binary_predict.py Dockerfile

Expected output:

Prediction: GOOD — Probabilities: good=0.998, bad=0.002

🗂 Extras

The full training and evaluation pipeline — including data preparation, training, validation, prediction — is available in the scripts/ folder.

💬 Note: Scripts are written with Polish comments and variable names for clarity during local development. Logic is fully portable.

📘 License

MIT

🙌 Credits

Model powered by Hugging Face Transformers
Tokenizer: CodeBERT
Rule definitions: Hadolint