Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3# Load model and tokenizer
4model_path = "./models/atomllama-33K-5x5-DigitMesh-sparse-q8"
5tokenizer = AutoTokenizer.from_pretrained(model_path)
6model = AutoModelForCausalLM.from_pretrained(
7 model_path,
8 dtype="auto",
9 device_map="auto"
10)
11
12# Example: Classify a 5x5 binary digit pattern (digit "0")
13pattern = "1 1 1 1 1 1 0 0 0 1 1 0 0 0 1 1 0 0 0 1 1 1 1 1 1"
14prompt = f"{pattern} <SEP>"
15
16# Tokenize and generate prediction
17inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
18inputs.pop("token_type_ids", None)
19
20outputs = model.generate(
21 **inputs,
22 max_new_tokens=1,
23 do_sample=False
24)
25
26# Decode the prediction
27prediction = tokenizer.decode(
28 outputs[0][len(inputs.input_ids[0]):],
29 skip_special_tokens=True
30).strip()
31
32print(f"Predicted digit: {prediction}") # Expected: "D0"1@misc{atomllama-33k-digitMesh-sparse-q8,
2 title={AtomLlama-33K-5x5-DigitMesh-Sparse-Q8: A 50% Sparse INT8 Quantized Model for Digit Recognition},
3 author={Jun Zhu},
4 year={2026},
5 howpublished={\url{https://huggingface.co/junzzhu/atomllama-33K-5x5-DigitMesh-sparse-q8}}
6}