# !pip install -U -q transformers accelerate sentencepiece bitsandbytes
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(
    "pszemraj/Yi-6B-200K-Llama-sharded", use_fast=False
)
model = AutoModelForCausalLM.from_pretrained(
    "pszemraj/Yi-6B-200K-Llama-sharded", load_in_4bit=True
)

prompt = "Custom non-commercial software licenses are just so fun, aren't they? The best thing about them is"


if torch.cuda.is_available():
    torch.cuda.empty_cache()

inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
tokens = model.generate(
    **inputs,
    max_new_tokens=192,
    temperature=0.75,
    top_p=0.9,
    no_repeat_ngram_size=4,
    do_sample=True,
    renormalize_logits=True,
    repetition_penalty=1.05,
)
print(tokenizer.decode(tokens[0], skip_special_tokens=True))