### 指示:
{instruction}

### 入力:
{input}

### 応答:
{response}

<s>[INST] {instruction} [/INST] {response}</s>

from datasets import Dataset

def create_instruction_dataset():
    """指示チューニング用データセットの作成"""
    data = [
        {
            "instruction": "以下の文章を要約してください。",
            "input": "人工知能（AI）は、機械が人間のような知能を持つ技術です。機械学習はAIの一分野で、データから学習するアルゴリズムを研究します。深層学習は、多層のニューラルネットワークを使った機械学習の手法です。",
            "output": "AIは機械に知能を持たせる技術で、機械学習はデータから学習し、深層学習は多層ニューラルネットワークを使用します。"
        },
        {
            "instruction": "以下の質問に答えてください。",
            "input": "日本で最も高い山は何ですか？",
            "output": "富士山です。標高は3,776メートルです。"
        },
        {
            "instruction": "以下の英文を日本語に翻訳してください。",
            "input": "Machine learning is a subset of artificial intelligence.",
            "output": "機械学習は人工知能の一分野です。"
        },
        {
            "instruction": "以下の文章の感情を「ポジティブ」「ネガティブ」「中立」で分類してください。",
            "input": "今日は素晴らしい天気で、とても気分がいいです。",
            "output": "ポジティブ"
        },
    ]

    return Dataset.from_list(data)

def format_instruction(sample):
    """サンプルをプロンプト形式に変換"""
    if sample["input"]:
        prompt = f"""### 指示:
{sample["instruction"]}

### 入力:
{sample["input"]}

### 応答:
{sample["output"]}"""
    else:
        prompt = f"""### 指示:
{sample["instruction"]}

### 応答:
{sample["output"]}"""

    return {"text": prompt}

# データセット作成
dataset = create_instruction_dataset()
dataset = dataset.map(format_instruction)

print("Sample data:")
print(dataset[0]["text"])

import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling
)

def prepare_model_and_tokenizer(model_name="gpt2"):
    """モデルとトークナイザーを準備"""
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    # パディングトークンの設定
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        model.config.pad_token_id = tokenizer.eos_token_id

    return model, tokenizer

def tokenize_function(examples, tokenizer, max_length=512):
    """テキストをトークン化"""
    return tokenizer(
        examples["text"],
        truncation=True,
        max_length=max_length,
        padding="max_length"
    )

def train_instruction_model(dataset, model, tokenizer, output_dir="./instruction_model"):
    """Instruction Tuningを実行"""
    # トークン化
    tokenized_dataset = dataset.map(
        lambda x: tokenize_function(x, tokenizer),
        remove_columns=dataset.column_names
    )

    # データコレーター
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False  # Causal LMなのでFalse
    )

    # 学習設定
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=3,
        per_device_train_batch_size=4,
        gradient_accumulation_steps=4,
        learning_rate=2e-5,
        warmup_steps=100,
        logging_steps=10,
        save_steps=500,
        fp16=torch.cuda.is_available(),
    )

    # トレーナー
    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=tokenized_dataset,
        data_collator=data_collator,
    )

    # 学習実行
    trainer.train()

    return model

# 使用例（実行にはGPUが必要）
# model, tokenizer = prepare_model_and_tokenizer("gpt2")
# trained_model = train_instruction_model(dataset, model, tokenizer)

def generate_response(model, tokenizer, instruction, input_text="", max_length=256):
    """学習済みモデルで応答を生成"""
    if input_text:
        prompt = f"""### 指示:
{instruction}

### 入力:
{input_text}

### 応答:
"""
    else:
        prompt = f"""### 指示:
{instruction}

### 応答:
"""

    inputs = tokenizer(prompt, return_tensors="pt")

    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"],
            max_length=max_length,
            num_return_sequences=1,
            temperature=0.7,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id
        )

    response = tokenizer.decode(outputs[0], skip_special_tokens=True)

    # 応答部分のみを抽出
    if "### 応答:" in response:
        response = response.split("### 応答:")[-1].strip()

    return response

# 使用例
# response = generate_response(trained_model, tokenizer, "猫について説明してください")
# print(response)

import torch
import torch.nn as nn

class InstructionTuningLoss(nn.Module):
    def __init__(self, ignore_index=-100):
        super().__init__()
        self.ignore_index = ignore_index
        self.loss_fn = nn.CrossEntropyLoss(ignore_index=ignore_index)

    def create_labels(self, input_ids, instruction_lengths):
        """応答部分のみをラベルとするマスクを作成"""
        labels = input_ids.clone()

        # 指示部分をマスク（ignore_indexに設定）
        for i, inst_len in enumerate(instruction_lengths):
            labels[i, :inst_len] = self.ignore_index

        return labels

    def forward(self, logits, input_ids, instruction_lengths):
        """損失を計算"""
        # ラベルを作成
        labels = self.create_labels(input_ids, instruction_lengths)

        # シフトして損失計算（次トークン予測）
        shift_logits = logits[..., :-1, :].contiguous()
        shift_labels = labels[..., 1:].contiguous()

        # 損失計算
        loss = self.loss_fn(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1)
        )

        return loss

# 使用例
# loss_fn = InstructionTuningLoss()
# loss = loss_fn(model_output.logits, input_ids, instruction_lengths)

def augment_instruction(instruction, input_text, output):
    """指示の言い換えによるデータ拡張"""
    variations = [
        f"次の指示に従ってください: {instruction}",
        f"以下のタスクを実行: {instruction}",
        f"あなたの仕事: {instruction}",
        instruction,  # オリジナル
    ]

    augmented = []
    for var in variations:
        augmented.append({
            "instruction": var,
            "input": input_text,
            "output": output
        })

    return augmented

def generate_instruction_data(llm_client, seed_tasks, n_generate=10):
    """Self-Instructによるデータ生成（疑似コード）"""
    prompt = f"""以下は指示と応答のペアの例です:

{format_seed_tasks(seed_tasks)}

同様の形式で、新しい指示と応答のペアを{n_generate}個生成してください:"""

    response = llm_client.generate(prompt)
    new_tasks = parse_generated_tasks(response)

    return new_tasks