#!/usr/bin/env python3
"""
Qwen3-8B 微调脚本
"""

import os
import sys
import torch

# 设置环境变量 - 必须在导入之前
os.environ['UNSLOTH_NO_STATISTICS'] = '1'  # 禁用Unsloth统计信息收集，减少内存占用
os.environ['HF_HUB_OFFLINE'] = '1'         # 强制使用本地缓存，避免从HuggingFace Hub下载
os.environ['TRANSFORMERS_OFFLINE'] = '1'   # 强制使用本地缓存，避免从HuggingFace Hub下载
os.environ['WANDB_MODE'] = 'disabled'      # 禁用Weights & Biases日志记录，避免登录问题

# 现在导入其他包
import json
from datasets import load_dataset
from unsloth import FastLanguageModel
from transformers import TrainingArguments
from trl import SFTTrainer

def format_conversation(example):
    """格式化对话数据 - 适配Qwen格式"""
    conversations = example["conversations"]
    
    formatted_text = ""
    for msg in conversations:
        if msg["role"] == "user":
            formatted_text += f"<|im_start|>user\n{msg['content']}<|im_end|>\n"
        elif msg["role"] == "assistant":
            formatted_text += f"<|im_start|>assistant\n{msg['content']}<|im_end|>\n"
    
    return {"text": formatted_text}

def main():
    print("=" * 60)
    print("Qwen3-8B 微调")
    print("=" * 60)
    
    # 模型参数
    model_path = "/home/yiyu/nfs/workspace/qwen3_08b_tunning/Qwen3-8B"
    train_file = "/home/yiyu/nfs/workspace/qwen3_08b_tunning/HUST-Student-Handbook/lora_hust_student_handbookt.jsonl"
    eval_file = ""
    output_dir = "/home/yiyu/nfs/workspace/qwen3_08b_tunning/qwen3-8b-finetuned"
    
    # 检查文件是否存在
    if not os.path.exists(model_path):
        print(f"错误: 模型路径不存在: {model_path}")
        sys.exit(1)
    
    if not os.path.exists(train_file):
        print(f"错误: 训练文件不存在: {train_file}")
        sys.exit(1)
    
    # 检查GPU
    print(f"\nGPU信息:")
    print(f"  CUDA可用: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"  GPU名称: {torch.cuda.get_device_name(0)}")
        print(f"  GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
    
    # 1. 加载模型和分词器
    print("\n1. 加载模型和分词器...")
    try:
        model, tokenizer = FastLanguageModel.from_pretrained(
            model_name=model_path,      # 预训练模型路径或HuggingFace模型ID
            max_seq_length=2048,        # 模型支持的最大序列长度
            dtype=None,                 # 数据类型，None表示自动检测（通常是torch.float16）
            load_in_4bit=True,          # 使用4位量化加载模型，减少内存占用
            token=None,                 # HuggingFace访问令牌，None表示不使用或已有缓存
        )
        print(f"  ✓ 模型加载成功")
    except Exception as e:
        print(f"  ✗ 模型加载失败: {e}")
        sys.exit(1)
    
    # 设置tokenizer
    tokenizer.pad_token = tokenizer.eos_token  # 使用EOS令牌作为填充令牌
    tokenizer.padding_side = "right"           # 在序列右侧进行填充
    print(f"  ✓ Tokenizer设置完成")
    
    # 2. 应用LoRA
    print("\n2. 应用LoRA配置...")
    try:
        model = FastLanguageModel.get_peft_model(
            model,                          # 基础模型
            r=16,                           # LoRA秩，控制低秩矩阵的维度，值越小参数越少
            target_modules=[                # 要应用LoRA的目标模块列表
                "q_proj", "k_proj", "v_proj", "o_proj",  # 注意力机制投影层
                "gate_proj", "up_proj", "down_proj",     # FFN层投影
            ],
            lora_alpha=32,                  # LoRA缩放因子，控制新权重对原始权重的贡献程度
            lora_dropout=0.05,              # LoRA层的Dropout率，防止过拟合
            bias="none",                    # 偏置训练策略："none"不训练，"all"训练所有偏置
            use_gradient_checkpointing=True, # 使用梯度检查点，减少内存占用但增加计算时间
            random_state=42,                # 随机种子，确保可重复性
            use_rslora=False,               # 是否使用rsLoRA（减少梯度的LoRA变体）
            loftq_config=None,              # LoftQ配置，用于量化感知微调
        )
        print(f"  ✓ LoRA应用成功 (r=16, alpha=32)")
    except Exception as e:
        print(f"  ✗ LoRA应用失败: {e}")
        sys.exit(1)
    
    # 计算可训练参数
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total_params = sum(p.numel() for p in model.parameters())
    print(f"  ✓ 可训练参数: {trainable_params:,} (占总参数 {100*trainable_params/total_params:.2f}%)")
    
    # 3. 加载数据
    print("\n3. 加载和准备数据...")
    try:
        # 加载训练数据
        train_dataset = load_dataset("json", data_files=train_file, split="train")
        train_dataset = train_dataset.map(
            format_conversation,
            remove_columns=["conversations"]
        )
        print(f"  ✓ 训练数据: {len(train_dataset)} 条样本")
        
        # 加载评估数据（如果有）
        if os.path.exists(eval_file):
            eval_dataset = load_dataset("json", data_files=eval_file, split="train")
            eval_dataset = eval_dataset.map(
                format_conversation,
                remove_columns=["conversations"]
            )
            print(f"  ✓ 评估数据: {len(eval_dataset)} 条样本")
            eval_data_available = True
        else:
            print(f"  ⓘ 评估数据不存在，跳过评估")
            eval_dataset = None
            eval_data_available = False
            
    except Exception as e:
        print(f"  ✗ 数据加载失败: {e}")
        sys.exit(1)
    
    # 4. 设置训练参数 - 使用兼容的参数名
    print("\n4. 设置训练参数...")
    try:
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 训练参数配置
        training_args = TrainingArguments(
            output_dir=output_dir,                    # 模型和日志输出目录
            num_train_epochs=3,                      # 训练轮数
            per_device_train_batch_size=2,           # 每个设备/GPU的训练批次大小
            gradient_accumulation_steps=4,           # 梯度累积步数，模拟更大批次
            warmup_ratio=0.03,                       # 学习率预热比例（占总训练步数的比例）
            learning_rate=2e-4,                      # 初始学习率
            fp16=True,                               # 使用混合精度训练（16位浮点数）
            logging_steps=10,                        # 每多少步记录一次日志
            save_strategy="steps",                   # 模型保存策略："steps"按步数保存
            save_steps=100,                          # 每多少步保存一次模型
            eval_strategy="steps" if eval_data_available else "no",  # 评估策略
            eval_steps=100 if eval_data_available else None,         # 每多少步评估一次
            gradient_checkpointing=True,             # 使用梯度检查点，减少内存占用
            optim="adamw_8bit",                      # 优化器类型，8位AdamW
            lr_scheduler_type="cosine",              # 学习率调度器类型：余弦退火
            seed=42,                                 # 随机种子
            report_to="none",                        # 禁用所有日志记录器
            ddp_find_unused_parameters=False,        # DDP训练中不查找未使用参数
            remove_unused_columns=False,             # 不自动删除数据集未使用列
            save_total_limit=3,                      # 最多保存的检查点数量
            load_best_model_at_end=True if eval_data_available else False,  # 训练结束时加载最佳模型
            metric_for_best_model="loss" if eval_data_available else None,  # 评估指标
            greater_is_better=False if eval_data_available else None,       # 指标是否越大越好
        )
        print(f"  ✓ 训练参数设置完成")
    except Exception as e:
        print(f"  ✗ 训练参数设置失败: {e}")
        sys.exit(1)
    
    # 5. 创建训练器
    print("\n5. 创建训练器...")
    try:
        trainer = SFTTrainer(
            model=model,                          # 要训练的模型
            tokenizer=tokenizer,                  # 分词器
            train_dataset=train_dataset,          # 训练数据集
            eval_dataset=eval_dataset,            # 评估数据集
            dataset_text_field="text",            # 数据集中文本字段的名称
            max_seq_length=2048,                  # 最大序列长度，超过部分会被截断
            packing=False,                        # 是否将多个短序列打包成一个批次
            args=training_args,                   # 训练参数配置
        )
        print(f"  ✓ 训练器创建成功")
    except Exception as e:
        print(f"  ✗ 训练器创建失败: {e}")
        sys.exit(1)
    
    # 6. 开始训练
    print("\n" + "=" * 60)
    print("开始训练...")
    print("=" * 60)
    
    try:
        # 训练模型
        train_result = trainer.train()
        
        print("\n" + "=" * 60)
        print("训练完成！")
        print("=" * 60)
        
        # 打印训练统计
        print(f"\n训练统计:")
        print(f"  总训练步数: {train_result.global_step}")
        print(f"  训练耗时: {train_result.metrics['train_runtime']:.2f} 秒")
        print(f"  每秒步数: {train_result.metrics['train_samples_per_second']:.2f}")
        
    except Exception as e:
        print(f"\n训练过程中出错: {e}")
        sys.exit(1)
    
    # 7. 保存模型
    print("\n6. 保存模型...")
    try:
        # 保存模型和tokenizer
        trainer.save_model()
        tokenizer.save_pretrained(output_dir)
        print(f"  ✓ 模型保存到: {output_dir}")
        
        # 保存训练参数
        with open(os.path.join(output_dir, "training_args.json"), "w") as f:
            json.dump(training_args.to_dict(), f, indent=2)
        print(f"  ✓ 训练参数保存")
        
    except Exception as e:
        print(f"  ✗ 模型保存失败: {e}")
        sys.exit(1)
    
    print("\n" + "=" * 60)
    print("微调流程完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
微调后的模型测试脚本 - 支持命令行参数输入模型路径
"""

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import sys
import os
import argparse


def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description="测试微调后的Qwen模型")
    parser.add_argument(
        "--model_path", 
        type=str, 
        default="/home/yiyu/nfs/workspace/qwen3_08b_tunning/Qwen3-8B",
        help="模型路径，默认为原始Qwen3-8B模型"
    )
    parser.add_argument(
        "--max_tokens",
        type=int,
        default=10000,
        help="最大生成token数，默认为10000"
    )
    
    # 解析命令行参数
    args = parser.parse_args()
    
    print(f"加载模型: {args.model_path}")
    print(f"最大生成token数: {args.max_tokens}")
    
    # 检查模型路径是否存在
    if not os.path.exists(args.model_path):
        print(f"错误: 模型路径不存在: {args.model_path}")
        sys.exit(1)
    
    # 加载tokenizer和模型
    try:
        # 加载tokenizer（分词器），用于文本编码和解码
        tokenizer = AutoTokenizer.from_pretrained(args.model_path)
        
        # 加载因果语言模型（用于文本生成）
        model = AutoModelForCausalLM.from_pretrained(
            args.model_path,  # 模型路径
            dtype=torch.float16,  # 模型权重数据类型，float16可减少内存使用，使用dtype而非torch_dtype
            device_map="auto",  # 自动将模型层分配到可用GPU上，支持多GPU
            trust_remote_code=True,  # 信任远程代码执行（对于某些自定义模型是必需的）
        )
    except Exception as e:
        print(f"加载模型时出错: {e}")
        sys.exit(1)
    
    # 测试问题
    test_questions = [
        "3年制研究生的毕业论文从开题到答辩通过最短需要多长时间？"
    ]
    
    print("\n" + "=" * 60)
    print("模型推理测试")
    print("=" * 60)
    
    for i, question in enumerate(test_questions, 1):
        print(f"\n问题 {i}: {question}")
        print("-" * 50)
        
        # 构建提示，使用Qwen模型的特定对话格式
        prompt = f"<|im_start|>user\n{question}<|im_end|>\n<|im_start|>assistant\n"
        
        # 编码输入文本为模型可理解的token ID
        inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
        
        # 生成文本
        try:
            with torch.no_grad():  # 禁用梯度计算，减少内存使用
                outputs = model.generate(
                    **inputs,  # 解包输入tensors（包含input_ids, attention_mask等）
                    max_new_tokens=args.max_tokens,  # 最大新生成的token数量（不包括输入）
                    temperature=0.7,  # 温度参数：控制随机性，较低的值使输出更确定性
                    top_p=0.9,  # 核采样参数：仅考虑累积概率达到top_p的token
                    do_sample=True,  # 启用采样模式（而非贪婪解码）
                    pad_token_id=tokenizer.eos_token_id,  # 指定填充token ID为结束token ID
                )
        except Exception as e:
            print(f"生成时出错: {e}")
            print("可能需要减少max_tokens值，或者检查GPU内存")
            continue
        
        # 解码生成的token ID为文本
        response = tokenizer.decode(outputs[0], skip_special_tokens=False)
        
        # 提取助手回复（仅保留assistant之后的文本）
        if "<|im_start|>assistant" in response:
            response = response.split("<|im_start|>assistant")[-1]
        if "<|im_end|>" in response:
            response = response.split("<|im_end|>")[0]
        
        print(f"回答: {response.strip()}")
        
        # 统计生成的token数量
        generated_tokens = outputs[0][inputs['input_ids'].shape[1]:]
        print(f"生成的token数: {len(generated_tokens)}")
    
    print("\n" + "=" * 60)
    print("测试完成！")
    print("=" * 60)


if __name__ == "__main__":
    main()

from nicegui import ui
import os
import glob
import time
from PIL import Image

# 自动管理图片路径和保存目录
INPUT_DIR = 'static/input'
DONE_DIR = 'static/done'
CROPS_DIR = 'static/crops'
GRAY_IMAGE = 'static/gray.png'

os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(DONE_DIR, exist_ok=True)
os.makedirs(CROPS_DIR, exist_ok=True)

DISPLAY_WIDTH = 1240
DISPLAY_HEIGHT = 1790
SELECT_BOX_HEIGHT = 300

select_box_y = 0
img = None
select_box = None

def ensure_gray_image():
    if not os.path.exists(GRAY_IMAGE):
        img_gray = Image.new('RGB', (DISPLAY_WIDTH, DISPLAY_HEIGHT), (200, 200, 200))
        img_gray.save(GRAY_IMAGE)

def get_input_images():
    return sorted(glob.glob(os.path.join(INPUT_DIR, '*.jpg')) + glob.glob(os.path.join(INPUT_DIR, '*.png')))

def get_first_image():
    images = get_input_images()
    if images:
        return images[0]
    ensure_gray_image()
    return GRAY_IMAGE

current_image_path = get_first_image()

def update_select_box(event):
    global select_box_y
    if event is None:
        y = select_box_y
    else:
        offset_y = getattr(event.args, 'offsetY', None)
        if offset_y is None:
            offset_y = event.args.get('offsetY', 0) if isinstance(event.args, dict) else 0
        y = max(0, min(offset_y, DISPLAY_HEIGHT - SELECT_BOX_HEIGHT))
    select_box_y = y
    select_box.style(f'top: {select_box_y}px; left: 0px; width: {DISPLAY_WIDTH}px; height: {SELECT_BOX_HEIGHT}px; position: absolute; border: 2px solid #00f; background: rgba(0,0,255,0.1); pointer-events: none;')

def crop_and_save(event):
    global select_box_y, current_image_path
    img_path = current_image_path
    if not os.path.exists(img_path):
        ui.notify('图片不存在')
        return
    try:
        img_obj = Image.open(img_path)
        img_width, img_height = img_obj.size
        scale = img_height / DISPLAY_HEIGHT
        crop_left = 0
        crop_upper = int(select_box_y * scale)
        crop_right = img_width
        crop_lower = int(min((select_box_y + SELECT_BOX_HEIGHT) * scale, img_height))
        crop_box = (crop_left, crop_upper, crop_right, crop_lower)
        crop_img = img_obj.crop(crop_box)
        if crop_img.mode == "RGBA":
            crop_img = crop_img.convert("RGB")
        ts = int(time.time() * 1000)
        save_path = os.path.join(CROPS_DIR, f'{ts}.jpg')
        crop_img.save(save_path)
        ui.notify(f'已保存: {save_path}')
    except Exception as e:
        ui.notify(f'裁剪失败: {e}')

def handle_mouse_click(event):
    global SELECT_BOX_HEIGHT
    button = event.args.get('button', 0)
    if button == 0:  # left
    #     SELECT_BOX_HEIGHT = max(10, SELECT_BOX_HEIGHT - 10)
    #     update_select_box(event)
    # elif button == 2:  # right
    #     SELECT_BOX_HEIGHT += 10
    #     update_select_box(event)
    # elif button == 1:  # middle
        crop_and_save(event)

def done_and_next():
    global current_image_path, img
    images = get_input_images()
    if current_image_path in images:
        basename = os.path.basename(current_image_path)
        # 检查done目录是否已存在同名文件，若存在则加时间戳
        done_path = os.path.join(DONE_DIR, basename)
        if os.path.exists(done_path):
            name, ext = os.path.splitext(basename)
            ts = int(time.time() * 1000)
            done_path = os.path.join(DONE_DIR, f"{name}_{ts}{ext}")
        os.rename(current_image_path, done_path)
        ui.notify(f'已移到: {done_path}')
    images = get_input_images()
    if images:
        current_image_path = images[0]
    else:
        ensure_gray_image()
        current_image_path = GRAY_IMAGE
    img.set_source(current_image_path)
    
def adjust_height(delta):
    global SELECT_BOX_HEIGHT
    if delta < 0:
        SELECT_BOX_HEIGHT = max(10, SELECT_BOX_HEIGHT + delta)
    else:
        SELECT_BOX_HEIGHT += delta
    update_select_box(None)

@ui.page('/')
def main_page():
    global img, select_box, current_image_path
    with ui.row().classes('w-full h-full justify-center items-center'):
        with ui.column().classes('w-full h-full items-center'):
            ui.label('审校区域标记工具').classes('text-2xl q-mb-md')
            ui.label('选择有审校修改的区域进行标记, 一页内容标记完成后点击“完成并下一张”按钮, 当图片显示为灰色时表示已经处理完')
            with ui.element('div').style(f'position: relative; width: {DISPLAY_WIDTH}px; height: {DISPLAY_HEIGHT}px;'):
                img = ui.image(current_image_path).style(f'width: {DISPLAY_WIDTH}px; height: {DISPLAY_HEIGHT}px; display: block; margin-left: auto; margin-right: auto;')
                select_box = ui.element('div').style(f'position: absolute; top: 0px; left: 0px; width: {DISPLAY_WIDTH}px; height: {SELECT_BOX_HEIGHT}px; border: 2px solid #00f; background: rgba(0,0,255,0.1); pointer-events: none;')
                ui.button('-', on_click=lambda: adjust_height(-10)).style('position: fixed; top: 10px; right: 10px; width: 30px; height: 30px;')
                ui.button('+', on_click=lambda: adjust_height(10)).style('position: fixed; top: 10px; right: 50px; width: 30px; height: 30px;')
                img.on('mousemove', update_select_box)
                img.on('mousemove', update_select_box)
                if current_image_path != GRAY_IMAGE:
                    img.on('mousedown', handle_mouse_click)
            # 添加一个与图片等宽、高度为100的按钮，点击后执行done_and_next
            ui.button('完成并下一张', on_click=done_and_next).style(f'width: {DISPLAY_WIDTH}px; height: 100px; font-size: 2rem; margin-top: 10px; margin-left: auto; margin-right: auto; display: block;')

ui.run(title='图片裁剪工具', native=False, port=8080)