import torch

# Define ByteLevelBPETokenizer and TransformerForCausalLM exactly as in the homework notebook.
# Then load artifacts from the Hub:

tokenizer = ByteLevelBPETokenizer.from_pretrained("Craftic/llm-course-hw1-small")
model = TransformerForCausalLM.from_pretrained("Craftic/llm-course-hw1-small")
model.eval()

prompt = "Муж приходит домой и говорит:"
input_ids = torch.tensor([tokenizer.encode(prompt)], dtype=torch.long)
attention_mask = torch.ones_like(input_ids)

with torch.no_grad():
    logits = model(input_ids, attention_mask)

next_token_id = logits[0, -1].argmax().item()
print(tokenizer.decode(tokenizer.encode(prompt, add_eos_token=False) + [next_token_id]))

Notes

This is an educational model, not a production-ready checkpoint.
Because the model is trained on jokes, generations may be low-quality, repetitive, or stylistically narrow.
The repository does not include a packaged inference library; the custom tokenizer/model classes should be copied from the homework notebook or moved into a Python module before loading.

Limitations

Small context window
Small vocabulary
Trained on a narrow-domain dataset
No safety alignment or moderation tuning

Intended use

Coursework demonstration
Experiments with custom tokenization and compact LMs
Lightweight local generation experiments

Not intended use

Factual QA
Safety-critical tasks
Production deployment without additional packaging, evaluation, and safeguards