1from transformers import AutoModelForCausalLM, AutoTokenizer
2
3model_name = "damfle/ornith-9b-custom"
4tokenizer = AutoTokenizer.from_pretrained(model_name)
5model = AutoModelForCausalLM.from_pretrained(
6 model_name,
7 device_map="auto",
8 load_in_4bit=True,
9 bnb_4bit_compute_dtype=torch.float16
10)
11
12input_text = "Your prompt here"
13inputs = tokenizer(input_text, return_tensors="pt").to("cuda")
14outputs = model.generate(**inputs, max_new_tokens=100)
15print(tokenizer.decode(outputs[0], skip_special_tokens=True))