Views
No views yet
1from peft import PeftModel
2from transformers import AutoModelForCausalLM, AutoTokenizer
3
4# Load the base model
5base_model = AutoModelForCausalLM.from_pretrained(
6 "meta-llama/Llama-3.2-3B-Instruct",
7 torch_dtype="auto",
8 device_map="auto"
9)
10tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")
11
12# Load the adapter
13model = PeftModel.from_pretrained(base_model, "atc_llama")
14
15# Process an ATC message
16instruction = "As an ATC communication expert, improve this transcript and analyze its intentions and data."
17message = "southwest five niner two turn left heading three four zero descend and maintain flight level two five zero"
18
19prompt = f"<|begin_of_text|><|header_start|>user<|header_end|>\n\n{instruction}\n\nOriginal: {message}<|eot|><|header_start|>assistant<|header_end|>\n\n"
20inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
21
22# Generate improved transcript and analysis
23outputs = model.generate(**inputs, max_new_tokens=512, do_sample=False)
24response = tokenizer.decode(outputs[0, inputs["input_ids"].shape[1]:], skip_special_tokens=True)
25print(response)