zerolimites's picture
BitNet b1.58 TinyDialogues: trained checkpoint, GGUF, scripts, patches
655f931 verified
Raw History Blame Contribute Delete
977 Bytes
#!/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))