Views
No views yet
1# First, make sure you have the necessary libraries installed:
2# pip install transformers peft bitsandbytes accelerate
3
4import torch
5from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
6from peft import PeftModel
7from accelerate import infer_auto_device_map, dispatch_model
8
9fine_tuned_model_id = "Falah/my-qlora-mistral7b-instruct"
10base_model_id = "mistralai/Mistral-7B-Instruct-v0.2"
11
12print("Loading tokenizer...")
13tokenizer = AutoTokenizer.from_pretrained(fine_tuned_model_id)
14
15print("Loading base model with quantization...")
16bnb_config = BitsAndBytesConfig(
17 load_in_4bit=True,
18 bnb_4bit_use_double_quant=True,
19 bnb_4bit_quant_type="nf4",
20 bnb_4bit_compute_dtype=torch.float16
21)
22base_model = AutoModelForCausalLM.from_pretrained(
23 base_model_id,
24 quantization_config=bnb_config,
25 device_map=None, # Load to CPU initially
26 torch_dtype=torch.float16,
27 trust_remote_code=True,
28)
29
30print("Loading PEFT adapter onto the base model...")
31model = PeftModel.from_pretrained(base_model, fine_tuned_model_id)
32
33print("Dispatching model to devices...")
34device_map = infer_auto_device_map(model, dtype=torch.float16)
35model = dispatch_model(model, device_map=device_map)
36
37# Ensure the model is in evaluation mode
38model.eval()
39
40print("Creating text generation pipeline...")
41generator = pipeline(
42 "text-generation",
43 model=model,
44 tokenizer=tokenizer,
45 torch_dtype=torch.float16,
46 device_map="auto",
47)
48
49# Define a sample user prompt
50user_prompt = "Write a short story about a robot learning to love."
51
52# Format the prompt
53formatted_prompt = f"[INST] {user_prompt} [/INST]"
54
55# Generate text
56outputs = generator(
57 formatted_prompt,
58 max_new_tokens=200,
59 num_return_sequences=1,
60 do_sample=True,
61 temperature=0.7,
62 top_k=50,
63 top_p=0.95,
64)
65
66# Print the generated text
67for i, output in enumerate(outputs):
68 print(f"Generated Output {i+1}:\n{output['generated_text']}")