Views
No views yet
1from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
2from peft import PeftModel, PeftConfig
3import torch
4import re
5
6# Load the adapter configuration
7config = PeftConfig.from_pretrained("SamuelJaja/llama-3.1-8b-instruct-construction-lora-a100")
8
9# Load base model with quantization
10bnb_config = BitsAndBytesConfig(load_in_8bit=True)
11model = AutoModelForCausalLM.from_pretrained(
12 config.base_model_name_or_path,
13 quantization_config=bnb_config,
14 device_map="auto"
15)
16
17# Load LoRA adapter
18model = PeftModel.from_pretrained(model, "SamuelJaja/llama-3.1-8b-instruct-construction-lora-a100")
19
20# Load tokenizer
21tokenizer = AutoTokenizer.from_pretrained(config.base_model_name_or_path)
22tokenizer.pad_token = tokenizer.eos_token
23
24# Clean response function
25def clean_response(text):
26 return re.sub(r'\[/?INST\]', '', text).strip()
27
28# Generate text
29def generate_response(prompt, temperature=0.1, max_tokens=256):
30 # Format properly
31 if not prompt.startswith("[INST]"):
32 formatted_prompt = f"[INST] {prompt} [/INST]"
33 else:
34 formatted_prompt = prompt
35
36 inputs = tokenizer(formatted_prompt, return_tensors="pt").to("cuda")
37
38 outputs = model.generate(
39 input_ids=inputs.input_ids,
40 attention_mask=inputs.attention_mask,
41 max_new_tokens=max_tokens,
42 temperature=temperature,
43 top_p=0.9,
44 do_sample=False
45 )
46
47 full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
48
49 # Remove prompt from output
50 if formatted_prompt in full_response:
51 response = full_response.replace(formatted_prompt, "").strip()
52 else:
53 response = full_response
54
55 # Clean any remaining instruction tags
56 response = clean_response(response)
57
58 return response
59
60# Example use
61question = "What are the main requirements for fire safety in commercial buildings?"
62answer = generate_response(question)
63print(answer)