Views
No views yet
en)1import os
2import torch
3from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
4from peft import PeftModel
5
6base_model_id = "google/gemma-2-2b-it"
7lora_adapter_repo = "saad9694/gemma-2-2b-it-noir-detective"
8
9# Configure 4-bit quantization for resource-constrained execution
10bnb_config = BitsAndBytesConfig(
11 load_in_4bit=True,
12 bnb_4bit_quant_type="nf4",
13 bnb_4bit_compute_dtype=torch.bfloat16
14)
15
16tokenizer = AutoTokenizer.from_pretrained(base_model_id)
17base_model = AutoModelForCausalLM.from_pretrained(
18 base_model_id,
19 quantization_config=bnb_config,
20 device_map="auto"
21)
22
23# Fetch and apply the LoRA adapter layers
24model = PeftModel.from_pretrained(base_model, lora_adapter_repo)
25model.eval()
26
27# Inference example
28prompt = "<start_of_turn>user\nWho killed the classical jazz musician?\n<end_of_turn>\n<start_of_turn>model\n"
29inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
30
31with torch.no_grad():
32 outputs = model.generate(**inputs, max_new_tokens=150, temperature=0.7, do_sample=True)
33
34print(tokenizer.decode(outputs[0][inputs['input_ids'].shape[1]:], skip_special_tokens=True))