File size: 4,293 Bytes
1b6fcc6 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 | #!/usr/bin/env python3
"""Predict whether Lean source code compiles with a trained classifier."""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from typing import Any, Sequence
DEFAULT_CACHE_DIR = None
DEFAULT_CLASSIFIER_OUTPUT = 'EPFL-LARA/lean-encoder-modernbert-large'
def encode_head_tail(
tokenizer: Any,
texts: Sequence[str],
*,
max_length: int,
head_fraction: float = 0.25,
) -> dict[str, torch.Tensor]:
"""Tokenize, retaining the beginning and end when a source is too long."""
if not 0.0 < head_fraction < 1.0:
raise ValueError("head_fraction must be between 0 and 1")
payload_budget = max_length - tokenizer.num_special_tokens_to_add(pair=False)
if payload_budget <= 0:
raise ValueError("max_length is smaller than the tokenizer's special tokens")
head_tokens = max(1, round(payload_budget * head_fraction))
tail_tokens = payload_budget - head_tokens
# A few compiler_data rows are over a million characters. Bound tokenizer
# memory with a generous character pre-slice, then perform the exact slice
# in token space below.
prepared_texts: list[str] = []
head_character_budget = max(2_048, head_tokens * 8)
tail_character_budget = max(2_048, tail_tokens * 8)
for text in texts:
if len(text) > head_character_budget + tail_character_budget:
text = (
text[:head_character_budget]
+ "\n-- middle omitted for classifier context window --\n"
+ text[-tail_character_budget:]
)
prepared_texts.append(text)
tokenized = tokenizer(
prepared_texts,
add_special_tokens=False,
truncation=False,
return_attention_mask=False,
verbose=False,
)["input_ids"]
features: list[dict[str, list[int]]] = []
for input_ids in tokenized:
if len(input_ids) > payload_budget:
input_ids = input_ids[:head_tokens] + input_ids[-tail_tokens:]
with_special_tokens = tokenizer.build_inputs_with_special_tokens(input_ids)
features.append(
{
"input_ids": with_special_tokens,
"attention_mask": [1] * len(with_special_tokens),
}
)
return tokenizer.pad(
features,
padding=True,
pad_to_multiple_of=8,
return_tensors="pt",
)
def main() -> None:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("path", nargs="?", help="Lean file; omit to read stdin")
parser.add_argument("--model", default=DEFAULT_CLASSIFIER_OUTPUT)
parser.add_argument("--cache-dir", default=DEFAULT_CACHE_DIR)
parser.add_argument("--max-length", type=int, default=8192)
parser.add_argument("--head-fraction", type=float, default=0.25)
args = parser.parse_args()
source_code = (
Path(args.path).read_text(encoding="utf-8") if args.path else sys.stdin.read()
)
if not source_code.strip():
parser.error("Lean source is empty")
tokenizer = AutoTokenizer.from_pretrained(args.model, cache_dir=args.cache_dir)
model = AutoModelForSequenceClassification.from_pretrained(
args.model,
cache_dir=args.cache_dir,
dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
attn_implementation="sdpa",
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device).eval()
inputs = encode_head_tail(
tokenizer,
[source_code],
max_length=args.max_length,
head_fraction=args.head_fraction,
)
inputs = {key: value.to(device) for key, value in inputs.items()}
with torch.inference_mode():
probabilities = model(**inputs).logits.float().softmax(dim=-1)[0].cpu()
valid_probability = float(probabilities[1])
print(
json.dumps(
{
"is_valid": valid_probability >= 0.5,
"valid_probability": valid_probability,
"invalid_probability": float(probabilities[0]),
},
indent=2,
)
)
if __name__ == "__main__":
main()
|