Views
No views yet
1pip install -q accelerate bitsandbytes trl datasets
2pip install git+https://github.com/huggingface/transformers
3pip install git+https://github.com/huggingface/peft.git1import os
2import torch
3from datasets import load_dataset
4from transformers import (
5 AutoModelForCausalLM,
6 AutoTokenizer,
7 BitsAndBytesConfig,
8 HfArgumentParser,
9 pipeline,
10 logging,
11)
12from peft import LoraConfig, PeftModel
13
14base_model_name = "mistralai/Mistral-7B-Instruct-v0.1"
15
16
17################################################################################
18# bitsandbytes parameters
19################################################################################
20
21# Activate 4-bit precision base model loading
22use_4bit = True
23
24# Compute dtype for 4-bit base models
25bnb_4bit_compute_dtype = "float16"
26
27# Quantization type (fp4 or nf4)
28bnb_4bit_quant_type = "nf4"
29
30# Activate nested quantization for 4-bit base models (double quantization)
31use_nested_quant = False
32
33# Load the entire model on the GPU 0
34device_map = {"": 0}
35
36# Load tokenizer and model with QLoRA configuration
37compute_dtype = getattr(torch, bnb_4bit_compute_dtype)
38
39bnb_config = BitsAndBytesConfig(
40 load_in_4bit=use_4bit,
41 bnb_4bit_quant_type=bnb_4bit_quant_type,
42 bnb_4bit_compute_dtype=compute_dtype,
43 bnb_4bit_use_double_quant=use_nested_quant,
44)
45
46# Check GPU compatibility with bfloat16
47if compute_dtype == torch.float16 and use_4bit:
48 major, _ = torch.cuda.get_device_capability()
49 if major >= 8:
50 print("=" * 80)
51 print("Your GPU supports bfloat16: accelerate training with bf16=True")
52 print("=" * 80)
53
54
55# Reload model in FP16 and merge it with LoRA weights
56base_model = AutoModelForCausalLM.from_pretrained(
57 base_model_name,
58 low_cpu_mem_usage=True,
59 return_dict=True,
60 torch_dtype=torch.float16,
61 quantization_config=bnb_config,
62 device_map=device_map,
63)
64model = PeftModel.from_pretrained(base_model, "Ashishkr/mistral-medical-consultation")
65model = model.merge_and_unload()
66
67# Reload tokenizer to save it
68tokenizer = AutoTokenizer.from_pretrained(base_model_name, trust_remote_code=True)
69tokenizer.pad_token = tokenizer.eos_token
70tokenizer.padding_side = "right"1# Run text generation pipeline with the merged model
2prompt = """
3i have a neck pain since 2 days .
4"""
5pipe = pipeline(task="text-generation", model=model, tokenizer=tokenizer,do_sample = True, temperature = 0.9, max_length=200)
6response = pipe(f"<s>[INST] {prompt} [/INST]")
7print(response[0]['generated_text'])
8