Views
No views yet
meta-llama/Llama-3.2-3B-Instructburgerbee/psychology_textbook1from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
2from peft import PeftModel
3import torch
4import warnings
5import logging
6from threading import Thread
7
8# Suppress warnings
9warnings.filterwarnings("ignore")
10
11# Suppress tqdm progress bar
12from transformers.utils import logging as hf_logging
13hf_logging.set_verbosity_error()
14
15# Suppress Accelerate and PEFT warnings
16logging.getLogger("peft").setLevel(logging.ERROR)
17logging.getLogger("bitsandbytes").setLevel(logging.ERROR)
18
19# Load base model
20device = 'cpu'
21model_name = "meta-llama/Llama-3.2-3B-Instruct"
22model = AutoModelForCausalLM.from_pretrained(
23 model_name,
24 device_map=device,
25 torch_dtype=torch.float16
26)
27
28# Load fine-tuned LoRA weights
29lora_model_name = "samuelchazy/fine-tuned-llama-3.2-3B-mlx-psychology"
30model = PeftModel.from_pretrained(model, lora_model_name)
31
32# Load tokenizer
33tokenizer = AutoTokenizer.from_pretrained(model_name)
34
35# Ensure padding and eos token are set
36tokenizer.pad_token_id = tokenizer.eos_token_id
37tokenizer.pad_token = tokenizer.eos_token
38tokenizer.padding_side = "left"
39
40print("Model loaded successfully!")
41
42# Define the prompt
43prompt = "What is the role of serotonin in mood regulation?"
44
45# Tokenize input
46inputs = tokenizer(prompt, return_tensors="pt").to(device)
47
48# Initialize the streamer
49streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
50
51# Launch streaming in a separate thread
52thread = Thread(target=model.generate, kwargs={
53 "inputs": inputs.input_ids,
54 "max_length": 768,
55 "streamer": streamer,
56 "do_sample": True,
57 "temperature": 0.5,
58})
59thread.start()
60
61# Read and stream output
62print("Generating response...\n")
63for token in streamer:
64 print(token, end="", flush=True)