Update Model_LOADER.py

Model_LOADER.py

@@ -1,210 +1,196 @@
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig, BitsAndBytesConfig,AutoConfig
-import time
-import torch
-torch.backends.cuda.matmul.allow_tf32 = True
-import random
-from datasets import load_dataset
-from transformers import TrainingArguments
-from trl import SFTTrainer
-from peft import LoraConfig
-# from accelerate import infer_auto_device_map, init_empty_weights, dispatch_model
-from torch.nn import CrossEntropyLoss
-torch.autograd.set_detect_anomaly(True)
-random_seed = 42
-torch.manual_seed(random_seed)
-random.seed(random_seed)
-# Set the device for each process
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-# torch.cuda.set_device(device)
-
-
-
-n_ahead_talk_global = 4
-n_passes_global = 2
-n_ahead_global = 8
-n_examples = 0
-
-def model_init(params):
-    original = False
-    if params is None:
-        params = {}
-    else:
-        params = params.params
-    # save params to file
-    n_ahead = params.get("n_ahead", n_ahead_global if not original else 1)
-    n_ahead_talk = params.get("n_ahead_talk", n_ahead_talk_global if not original else 1)
-    n_passes = params.get("n_passes", n_passes_global if not original else 1)
-    gumbel_temperature = params.get("gumbel_temperature", 1)
-    use_start_thought_token = params.get("use_start_thought_token", True)
-    use_end_thought_token = params.get("use_end_thought_token", True)
-    include_policy_loss = params.get("include_policy_loss", True)
-    gumbel_detach = params.get("gumbel_detach", True)
-    merged_talk_heads = params.get("merged_talk_heads", True)
-    residual_think_head = params.get("residual_think_head", False)
-    optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)
-
-    model_id = "LeroyDyer/_Spydaz_Web_AI_V2_Aligned"
-    tokenizer_id = model_id
-    print("Loading model")
-
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
-        max_thoughts=n_ahead + n_ahead_talk + 1,
-        merged_talk_heads=merged_talk_heads,
-        merged_lm_and_talk_heads=False,
-        merged_lm_and_think_heads=True,
-        use_concat_talk_head=True,
-        use_shallow_think=True,
-        use_shallow_talk=False,
-        use_complex_think_head=False,
-        use_complex_talk_head=True,
-        use_weighted_talk_head=True,
-        trust_remote_code=True,
-        device_map="auto",
-    )
-    print("Loaded model")
-
-    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, truncation=True, padding_side="right")
-    tokenizer.pad_token_id = tokenizer.eos_token_id
-
-    special_tokens_to_add = []
-    if model.use_start_thought_token:
-        special_tokens_to_add.append("<|startthought|>")
-    if model.use_end_thought_token:
-        special_tokens_to_add.append("<|endthought|>")
-    if special_tokens_to_add:
-        tokenizer.add_special_tokens({"additional_special_tokens": special_tokens_to_add})
-        model.resize_token_embeddings(len(tokenizer))
-    model.tokenizer = tokenizer
-    for name, module in model.named_modules():
-        if "embed" in name:
-            print(module, flush=True)
-
-    model.gumbel_detach = gumbel_detach
-    model.include_policy_loss = include_policy_loss
-    model.use_end_thought_token = use_end_thought_token
-    model.use_start_thought_token = use_start_thought_token
-    model.n_ahead = n_ahead
-    model.n_ahead_talk = n_ahead_talk
-    model.n_passes = n_passes
-    model.residual_think_head = residual_think_head
-    model.optimize_lm_head_only_at_start = optimize_lm_head_only_at_start
-    model.gumbel_temperature = gumbel_temperature
-    model.original_mode = original
-    model.config_params = params
-    return model
-model,tokenizer = model_init(None)
-tokenizer.save_pretrained("IpretrainedModel")
-model.save_pretrained("IpretrainedModel")
-
-
-## TRAINING : 
-
-peft_config = LoraConfig(
-          r = 128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
-    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
-                      "gate_proj", "up_proj", "down_proj","lm_head", "embed_tokens"],
-    lora_alpha = 32,
-    lora_dropout = 0, # Supports any, but = 0 is optimized
-    bias = "none", 
-    use_dora=True,
-)
-
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig
-from datasets import load_dataset
-from transformers import TrainingArguments
-from trl import SFTTrainer
-from peft import LoraConfig
-
-## DATA
-alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
-
-### Instruction:
-{}
-
-### Input:
-{}
-
-### Response:
-{}"""
-EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
-def formatting_prompts_func(examples):
-    instructions = examples["instruction"]
-    inputs       = examples["input"]
-    outputs      = examples["output"]
-    texts = []
-    for instruction, input, output in zip(instructions, inputs, outputs):
-        # Must add EOS_TOKEN, otherwise your generation will go on forever!
-        text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
-        texts.append(text)
-    return { "text" : texts, }
-pass
-dataset = load_dataset("gate369/Alpaca-Star", split = "train[:2000]")
-dataset = dataset.shuffle(seed=3704)
-dataset = dataset.map(formatting_prompts_func, batched = True,)
-## TRAIN
-model
-model.train
-max_seq_length = 32000
-training_args = TrainingArguments(
-    output_dir="./out",
-    num_train_epochs=3,
-    per_device_train_batch_size=1,
-    gradient_checkpointing=False,
-    gradient_accumulation_steps=8,
-    optim="lion_32bit",
-    logging_steps=1,
-    save_strategy="steps",
-    save_steps=300,
-    max_steps=1000,
-    bf16=True,
-    tf32=False,
-    learning_rate=6e-05,
-    max_grad_norm=0.3,
-    warmup_ratio=0.06,
-    lr_scheduler_type="cosine",
-    push_to_hub=False,
-
-)
-trainer = SFTTrainer(
-    args=training_args,
-    train_dataset=dataset,
-    model=model,
-    tokenizer=tokenizer,
-    max_seq_length=max_seq_length,
-    dataset_text_field="text",
-    peft_config=peft_config,
-)
-trainer.train()
-
-## SAVE
-tokenizer.save_pretrained("SFTTrainerModel")
-model.save_pretrained("SFTTrainerModel")
-
-
-import os
-import huggingface_hub
-from huggingface_hub import notebook_login
-from huggingface_hub import create_repo, HfApi
-from huggingface_hub import hf_hub_download
-from huggingface_hub import create_repo, HfApi
-from huggingface_hub import snapshot_download
-
-MODEL_NAME = "_Spydaz_Web_AI_MistralStar"
-Folderinput = "SFTTrainerModel"
-WRITE_TOKEN = ""
-username = "LeroyDyer"
-huggingface_hub.login(WRITE_TOKEN)
-api = HfApi(token=WRITE_TOKEN)
-# Create empty repo
-api.create_repo(
-    repo_id = f"{username}/{MODEL_NAME}",
-    repo_type="model",
-    exist_ok=True,
-)
-
-api.upload_folder(
-    repo_id = f"{username}/{MODEL_NAME}",
-    folder_path = Folderinput
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig, BitsAndBytesConfig,AutoConfig
+import time
+import torch
+torch.backends.cuda.matmul.allow_tf32 = True
+import random
+from datasets import load_dataset
+from transformers import TrainingArguments
+from trl import SFTTrainer
+from peft import LoraConfig
+# from accelerate import infer_auto_device_map, init_empty_weights, dispatch_model
+from torch.nn import CrossEntropyLoss
+torch.autograd.set_detect_anomaly(True)
+random_seed = 42
+torch.manual_seed(random_seed)
+random.seed(random_seed)
+# Set the device for each process
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# torch.cuda.set_device(device)
+
+
+
+n_ahead_talk_global = 4
+n_passes_global = 2
+n_ahead_global = 8
+n_examples = 0
+
+def model_init(params):
+    original = False
+    if params is None:
+        params = {}
+    else:
+        params = params.params
+    # save params to file
+    n_ahead = params.get("n_ahead", n_ahead_global if not original else 1)
+    n_ahead_talk = params.get("n_ahead_talk", n_ahead_talk_global if not original else 1)
+    n_passes = params.get("n_passes", n_passes_global if not original else 1)
+    gumbel_temperature = params.get("gumbel_temperature", 1)
+    use_start_thought_token = params.get("use_start_thought_token", True)
+    use_end_thought_token = params.get("use_end_thought_token", True)
+    include_policy_loss = params.get("include_policy_loss", True)
+    gumbel_detach = params.get("gumbel_detach", True)
+    merged_talk_heads = params.get("merged_talk_heads", True)
+    residual_think_head = params.get("residual_think_head", False)
+    optimize_lm_head_only_at_start = params.get("optimize_lm_head_only_at_start", False)
+
+    model_id = "LeroyDyer/SpydazWeb_AGI_MistralStar"
+    tokenizer_id = model_id
+    print("Loading model")
+
+    model = AutoModelForCausalLM.from_pretrained(
+        model_id,
+        torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
+        max_thoughts=n_ahead + n_ahead_talk + 1,
+        merged_talk_heads=merged_talk_heads,
+        merged_lm_and_talk_heads=False,
+        merged_lm_and_think_heads=True,
+        use_concat_talk_head=True,
+        use_shallow_think=True,
+        use_shallow_talk=False,
+        use_complex_think_head=False,
+        use_complex_talk_head=True,
+        use_weighted_talk_head=True,
+        trust_remote_code=True,
+        device_map="auto",
+    )
+    print("Loaded model")
+
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_id, truncation=True, padding_side="right")
+    tokenizer.pad_token_id = tokenizer.eos_token_id
+
+
+
+    model.gumbel_detach = gumbel_detach
+    model.include_policy_loss = include_policy_loss
+    model.use_end_thought_token = use_end_thought_token
+    model.use_start_thought_token = use_start_thought_token
+    model.n_ahead = n_ahead
+    model.n_ahead_talk = n_ahead_talk
+    model.n_passes = n_passes
+    model.residual_think_head = residual_think_head
+    model.optimize_lm_head_only_at_start = optimize_lm_head_only_at_start
+    model.gumbel_temperature = gumbel_temperature
+    model.original_mode = original
+    model.config_params = params
+    return model
+    
+model,tokenizer = model_init(None)
+
+
+## TRAINING : 
+
+peft_config = LoraConfig(
+          r = 128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules = ["q_proj", "k_proj", "v_proj", "o_proj",
+                      "gate_proj", "up_proj", "down_proj","lm_head", "embed_tokens"],
+    lora_alpha = 32,
+    lora_dropout = 0, # Supports any, but = 0 is optimized
+    bias = "none", 
+    use_dora=True,
+)
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextGenerationPipeline, AutoConfig
+from datasets import load_dataset
+from transformers import TrainingArguments
+from trl import SFTTrainer
+from peft import LoraConfig
+
+## DATA
+alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.
+
+### Instruction:
+{}
+
+### Input:
+{}
+
+### Response:
+{}"""
+EOS_TOKEN = tokenizer.eos_token # Must add EOS_TOKEN
+def formatting_prompts_func(examples):
+    instructions = examples["instruction"]
+    inputs       = examples["input"]
+    outputs      = examples["output"]
+    texts = []
+    for instruction, input, output in zip(instructions, inputs, outputs):
+        # Must add EOS_TOKEN, otherwise your generation will go on forever!
+        text = alpaca_prompt.format(instruction, input, output) + EOS_TOKEN
+        texts.append(text)
+    return { "text" : texts, }
+pass
+dataset = load_dataset("gate369/Alpaca-Star", split = "train[:2000]")
+dataset = dataset.shuffle(seed=3704)
+dataset = dataset.map(formatting_prompts_func, batched = True,)
+## TRAIN
+
+max_seq_length = 1024
+training_args = TrainingArguments(
+    output_dir="./out",
+    num_train_epochs=3,
+    per_device_train_batch_size=1,
+    gradient_checkpointing=False,
+    gradient_accumulation_steps=8,
+    optim="lion_32bit",
+    logging_steps=1,
+    save_strategy="steps",
+    max_steps=1000,
+    bf16=True,
+    tf32=False,
+    learning_rate=6e-05,
+    max_grad_norm=0.3,
+    warmup_ratio=0.06,
+    lr_scheduler_type="cosine",
+    push_to_hub=False,
+
+)
+trainer = SFTTrainer(
+    args=training_args,
+    train_dataset=dataset,
+    model=model,
+    tokenizer=tokenizer,
+    max_seq_length=max_seq_length,
+    dataset_text_field="text",
+    peft_config=peft_config,
+)
+trainer.train()
+
+## SAVE
+tokenizer.save_pretrained("SFTTrainerModel")
+model.save_pretrained("SFTTrainerModel")
+
+
+import os
+import huggingface_hub
+from huggingface_hub import notebook_login
+from huggingface_hub import create_repo, HfApi
+from huggingface_hub import hf_hub_download
+from huggingface_hub import create_repo, HfApi
+from huggingface_hub import snapshot_download
+
+MODEL_NAME = "_Spydaz_Web_AI_MistralStar"
+Folderinput = "SFTTrainerModel"
+WRITE_TOKEN = ""
+username = "LeroyDyer"
+huggingface_hub.login(WRITE_TOKEN)
+api = HfApi(token=WRITE_TOKEN)
+# Create empty repo
+api.create_repo(
+    repo_id = f"{username}/{MODEL_NAME}",
+    repo_type="model",
+    exist_ok=True,
+)
+
+api.upload_folder(
+    repo_id = f"{username}/{MODEL_NAME}",
+    folder_path = Folderinput
 )