Module: Trainers::SaveUtils

Defined in:
lib/trainers/save_utils.rb

Class Method Summary collapse

Class Method Details

.load_lora_adapters(model, input_dir) ⇒ Object 



64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
 
# File 'lib/trainers/save_utils.rb', line 64

def self.load_lora_adapters(model, input_dir)
  adapter_path = File.join(input_dir, "adapter_model.safetensors")
  raise "Adapter file not found: #{adapter_path}" unless File.exist?(adapter_path)

  require "safetensors"
  lora_state = Safetensors::Torch.load_file(adapter_path)

  model.named_modules.each do |name, mod|
    if mod.is_a?(Trainers::LoraLinear)
      a_key = "#{name}.lora_A"
      b_key = "#{name}.lora_B"
      if lora_state[a_key] && lora_state[b_key]
        mod.load_lora_weights(lora_state[a_key], lora_state[b_key])
      end
    end
  end

  puts "LoRA adapters loaded from #{adapter_path}"
end

.save_lora_adapters(model, output_dir) ⇒ Object 



39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
 
# File 'lib/trainers/save_utils.rb', line 39

def self.save_lora_adapters(model, output_dir)
  FileUtils.mkdir_p(output_dir)

  lora_state = {}
  model.named_modules.each do |name, mod|
    if mod.is_a?(Trainers::LoraLinear)
      mod.lora_state_dict.each do |param_name, tensor|
        lora_state["#{name}.#{param_name}"] = tensor
      end
    end
  end

  if lora_state.empty?
    puts "No LoRA adapters found in model"
    return
  end

  adapter_path = File.join(output_dir, "adapter_model.safetensors")
  require "safetensors"
  Safetensors::Torch.save_file(lora_state, adapter_path)
  puts "LoRA adapters saved to #{adapter_path} (#{lora_state.size} tensors)"

  output_dir
end

.save_pretrained(model, tokenizer, output_dir, training_args: nil) ⇒ Object 



8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
 
# File 'lib/trainers/save_utils.rb', line 8

def self.save_pretrained(model, tokenizer, output_dir, training_args: nil)
  FileUtils.mkdir_p(output_dir)

  # Save model weights via safetensors
  state_dict = model.state_dict
  tensors = {}
  state_dict.each do |name, param|
    tensors[name] = param.is_a?(Torch::NN::Parameter) ? param.data : param
  end

  model_path = File.join(output_dir, "model.safetensors")
  require "safetensors"
  Safetensors::Torch.save_file(tensors, model_path)
  puts "Model saved to #{model_path}"

  # Save tokenizer
  if tokenizer
    if tokenizer.respond_to?(:save)
      tokenizer.save(File.join(output_dir, "tokenizer.json"))
    end
  end

  # Save training args
  if training_args
    args_path = File.join(output_dir, "training_args.json")
    File.write(args_path, JSON.pretty_generate(training_args.to_h))
  end

  output_dir
end