PardisSzah/PEFT_TuringQ_llama3_FT

from torch.utils.data import Dataset, DataLoader
from time import time
from datasets import load_dataset
from peft import LoraConfig, PeftModel, prepare_model_for_kbit_training, PeftConfig
from transformers import (
    AutoConfig,
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    TrainingArguments,
    pipeline,
)
from trl import SFTTrainer, setup_chat_format
from huggingface_hub import login
from tqdm import tqdm
import pandas as pd
import torch
import os

# Load the base model
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=True
)
base_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    trust_remote_code=True
)
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_id)

# Set up chat format
base_model, tokenizer = setup_chat_format(base_model, tokenizer)

# Load the PEFT adapter
peft_model_id = "PardisSzah/PEFT_TuringQ_llama3_FT"
peft_config = PeftConfig.from_pretrained(peft_model_id)
model = PeftModel.from_pretrained(base_model, peft_model_id)

# Merge the base model and PEFT adapter
model = model.merge_and_unload()

# Set the model to evaluation mode
model.eval()

pipeline = transformers.pipeline(
    "text-generation",
    model=model,
    tokenizer=tokenizer,
    torch_dtype=torch.float16,
    device_map="auto",
)

# You can change the parameters to match your specifications
def get_answer(system_message, user_message, temperature=0.8, max_length=1024):
    user_message = "Question: " + user_message + " Answer:"
    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_message},
    ]
    prompt = pipeline.tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    terminators = [
        pipeline.tokenizer.eos_token_id,
        pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>")
    ]
    sequences = pipeline(
        prompt,
        do_sample=True,
        top_p=0.9,
        temperature=temperature,
        eos_token_id=terminators,
        max_new_tokens=max_length,
        return_full_text=False,
        pad_token_id=pipeline.model.config.eos_token_id
    )
    answer = sequences[0]['generated_text']
    return answer 

# Use your own prompt
system_message =
"""
You are an AI assistant designed to answer theory of computation and math questions.
Please provide a step-by-step solution to the problem.
"""

answer = query_model(system_message, question, temperature=0.1, max_length=600)
print(answer)