llm_competition_2024_gemma2_9b_LoRA

Generate

次のコードを動かして実行(データセットの名前は適宜変更)

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
    HfArgumentParser,
    TrainingArguments,
    pipeline,
    logging,
)
from peft import (
    LoraConfig,
    PeftModel,
    prepare_model_for_kbit_training,
    get_peft_model,
)
import os, torch
from datasets import load_dataset
from trl import SFTTrainer, setup_chat_format

from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    BitsAndBytesConfig,
)
import torch
from tqdm import tqdm
import json
model_name="ruuuin/llm_competition_2024_gemma2_9b_LoRA"
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=False,
)
# Load model
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto",
)

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True,)
datasets = []
with open("datasets.json", "r") as f:
    item = ""
    for line in f:
      line = line.strip()
      item += line
      if item.endswith("}"):
        datasets.append(json.loads(item))
        item = ""
# gemma
results = []
for data in tqdm(datasets):

  input = data["input"]
  prompt = f"""### 指示
  {input}
  ### 回答：
  """

  input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
  outputs = model.generate(**input_ids, max_new_tokens=512, do_sample=False, repetition_penalty=1.2,)
  output = tokenizer.decode(outputs[0][input_ids.input_ids.size(1):], skip_special_tokens=True)

  results.append({"task_id": data["task_id"], "input": input, "output": output})