import json

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from predict import encode_head_tail

MODEL_ID = 'EPFL-LARA/lean-encoder-modernbert-large'
source_code = """import Mathlib

theorem nat_add_zero_identity (n : ℕ) : n + 0 = n := by
  rfl
"""

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_ID, dtype=dtype, attn_implementation="sdpa",
).to(device).eval()

inputs = encode_head_tail(
    tokenizer, [source_code], max_length=8192, head_fraction=0.25,
)
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.inference_mode():
    probabilities = model(**inputs).logits[0].float().softmax(dim=-1)

probability_valid = float(probabilities[1])  # Class 0: INVALID; class 1: VALID.
print(json.dumps({
    "is_valid": probability_valid >= 0.5,
    "valid_probability": probability_valid,
    "invalid_probability": float(probabilities[0]),
}, indent=2))