#!/usr/bin/env python3
"""HF demo with ternary BitLinear forward (W1.58A8) on CPU."""
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import BitNetForCausalLM, LlamaTokenizer
from train_bitnet import bitnetize

ap = argparse.ArgumentParser()
ap.add_argument('--model', default='/home/user/bitnet-tinydialogues')
ap.add_argument('--prompt', default='User: Hello Assistant:')
ap.add_argument('--tokens', type=int, default=64)
args = ap.parse_args()

torch.set_num_threads(2)
tok = LlamaTokenizer.from_pretrained(args.model, legacy=True)
model = bitnetize(BitNetForCausalLM.from_pretrained(args.model)).eval()
ids = tok(args.prompt, return_tensors='pt').input_ids
with torch.no_grad():
    out = model.generate(ids, max_new_tokens=args.tokens, do_sample=True, temperature=0.8,
                         top_p=0.9, eos_token_id=tok.eos_token_id, pad_token_id=tok.pad_token_id)
print(tok.decode(out[0], skip_special_tokens=True))