jameslloyd2/

ARTCAST4D-Chatbot

import torch  
from peft import PeftModel  
from transformers import (  
    AutoModelForCausalLM,  
    AutoTokenizer,  
    BitsAndBytesConfig,  
)
  
# Define paths and tokens for model loading
MODEL_PATH = "mistralai/Mistral-7B-Instruct-v0.2"  # Base model from Hugging Face
PEFT_PATH = "./model_objects/"  # Path to your fine-tuned LoRA/QLoRA adapters
HF_TOKEN = "hugging-face-token"  # Your Hugging Face authentication token
  
  
def load_model():  
    """    
    Loads a pre-trained language model with specified quantization configuration and additional optimisations.    
    Returns:        model: A pre-trained language model with quantization and other optimisations applied.
    """  
    # Set compute dtype to float16 for better performance and memory efficiency
    compute_dtype = getattr(torch, "float16")  
    
    # Configure 4-bit quantization using BitsAndBytesConfig
    bnb_config = BitsAndBytesConfig(  
        load_in_4bit=True,  # Enable 4-bit quantization to reduce memory usage
        bnb_4bit_quant_type="nf4",  # Use NormalFloat4 quantization method
        bnb_4bit_compute_dtype=compute_dtype,  # Set computation dtype to float16
        bnb_4bit_use_double_quant=False,  # Disable double quantization (saves memory but may reduce quality)
    )  
  
    # Load the base model with quantization configuration
    model = AutoModelForCausalLM.from_pretrained(  
        MODEL_PATH,  # Path to the base Mistral model
        quantization_config=bnb_config,  # Apply the quantization config
        device_map="cuda",  # Load model on GPU(s)
        token=HF_TOKEN,  # Authentication token for Hugging Face
    )  
  
    # Load and apply the fine-tuned PEFT (LoRA/QLoRA) adapters
    model = PeftModel.from_pretrained(model_id=PEFT_PATH, model=model)
    
    # Merge the adapters with the base model and unload adapters to save memory
    model = model.merge_and_unload()  
  
    return model  
  
  
def load_tokenizer():  
    """    
    Loads a pre-trained tokenizer for the model.    
    Returns:        tokenizer: A pre-trained tokenizer.
    """  
    # Load tokenizer from config
    tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, token=HF_TOKEN)
    
    # Set padding token to end-of-sequence token (required for some models)
    tokenizer.pad_token = tokenizer.eos_token  
    
    # Set padding side to "right" (tokens added to the right of the sequence)
    tokenizer.padding_side = "right"  
  
    return tokenizer  
  
# Load the fine-tuned model and tokenizer
model = load_model()
tokenizer = load_tokenizer()

# Define input text for generation
text = "Can I ask a question about disinformation?"  
  
# Tokenize the input text and convert to PyTorch tensors
batch = tokenizer(text, return_tensors='pt')  
  
# Generate text using automatic mixed precision for better performance
with torch.cuda.amp.autocast():  
    output_tokens = model.generate(
        **batch,  # Pass tokenized input
        max_new_tokens=1000,  # Maximum number of new tokens to generate
        do_sample=True  # Use sampling for more diverse outputs (vs greedy decoding)
    )  

# Decode the generated tokens back to text and print the result
print('\n\n', tokenizer.decode(output_tokens[0], skip_special_tokens=True))