Upload 17 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+stage1_simcse/epoch1/tokenizer.json filter=lfs diff=lfs merge=lfs -text
+stage1_simcse/final/tokenizer.json filter=lfs diff=lfs merge=lfs -text

stage1_simcse/epoch1/README.md

@@ -0,0 +1,202 @@
+---
+base_model: google/gemma-3-1b-pt
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.13.1

stage1_simcse/epoch1/adapter_config.json

@@ -0,0 +1,29 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "google/gemma-3-1b-pt",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "q_proj",
+    "v_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

stage1_simcse/epoch1/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4adbdf429104dd37e9287bd2c7d71cc99456cf20c480e60714fd9260dbaab71a
+size 2995512

stage1_simcse/epoch1/added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<image_soft_token>": 262144
+}

stage1_simcse/epoch1/special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "boi_token": "<start_of_image>",
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eoi_token": "<end_of_image>",
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "image_token": "<image_soft_token>",
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

stage1_simcse/epoch1/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4667f2089529e8e7657cfb6d1c19910ae71ff5f28aa7ab2ff2763330affad795
+size 33384568

stage1_simcse/epoch1/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c
+size 4689074

stage1_simcse/final/README.md

@@ -0,0 +1,202 @@
+---
+base_model: google/gemma-3-1b-pt
+library_name: peft
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]
+### Framework versions
+
+- PEFT 0.13.1

stage1_simcse/final/adapter_config.json

@@ -0,0 +1,29 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "google/gemma-3-1b-pt",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 16,
+  "lora_dropout": 0.0,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 8,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": [
+    "q_proj",
+    "v_proj"
+  ],
+  "task_type": "CAUSAL_LM",
+  "use_dora": false,
+  "use_rslora": false
+}

stage1_simcse/final/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4adbdf429104dd37e9287bd2c7d71cc99456cf20c480e60714fd9260dbaab71a
+size 2995512

stage1_simcse/final/added_tokens.json

@@ -0,0 +1,3 @@
+{
+  "<image_soft_token>": 262144
+}

stage1_simcse/final/special_tokens_map.json

@@ -0,0 +1,33 @@
+{
+  "boi_token": "<start_of_image>",
+  "bos_token": {
+    "content": "<bos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eoi_token": "<end_of_image>",
+  "eos_token": {
+    "content": "<eos>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "image_token": "<image_soft_token>",
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

stage1_simcse/final/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4667f2089529e8e7657cfb6d1c19910ae71ff5f28aa7ab2ff2763330affad795
+size 33384568

stage1_simcse/final/tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1299c11d7cf632ef3b4e11937501358ada021bbdf7c47638d13c0ee982f2e79c
+size 4689074

train1.py

@@ -0,0 +1,220 @@
+#!/usr/bin/env python3
+import os
+os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
+import sys
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from transformers import (
+    AutoTokenizer,
+    AutoModelForCausalLM,
+    get_linear_schedule_with_warmup
+)
+from peft import LoraConfig, get_peft_model, TaskType
+from datasets import load_dataset
+from tqdm.auto import tqdm
+from multiprocessing import freeze_support
+
+def main():
+    # Config
+    MODEL_NAME     = "google/gemma-3-1b-pt"
+    DATA_FILE      = "text.txt"       # one sequence per line
+    BATCH_SIZE     = 12
+    MAX_LENGTH     = 128
+    LR             = 1e-5
+    WEIGHT_DECAY   = 0.01
+    NUM_EPOCHS     = 1
+    VAL_RATIO      = 0.1              # 10% for validation
+    LORA_R         = 8
+    LORA_ALPHA     = 16
+    LORA_DROPOUT   = 0.0
+    PROJ_HIDDEN    = 512              # intermediate MLP width
+    TEMP           = 0.05
+    OUTPUT_DIR     = "stage1_simcse"
+    GRAD_CLIP_NORM = 1.0
+    SIM_CLAMP_MIN  = -10.0
+    SIM_CLAMP_MAX  = 10.0
+    SEED           = 42
+
+    os.makedirs(OUTPUT_DIR, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # enable TF32 and cuDNN autotuner on CUDA
+    if device.type == "cuda":
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32   = True
+        torch.backends.cudnn.benchmark     = True
+
+    # tokenizer + model
+    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+    base_model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        attn_implementation="eager"
+    )
+
+    # LoRA on q_proj & v_proj
+    lora_cfg = LoraConfig(
+        task_type=TaskType.CAUSAL_LM,
+        inference_mode=False,
+        r=LORA_R,
+        lora_alpha=LORA_ALPHA,
+        lora_dropout=LORA_DROPOUT,
+        target_modules=["q_proj", "v_proj"],
+    )
+    model_lora = get_peft_model(base_model, lora_cfg)
+
+    # Encoder + projection head now outputs hidden_size
+    class GemmaSimCSE(nn.Module):
+        def __init__(self, base):
+            super().__init__()
+            self.base = base
+            hs = base.config.hidden_size
+            self.proj = nn.Sequential(
+                nn.Linear(hs, PROJ_HIDDEN),
+                nn.ReLU(),
+                nn.Linear(PROJ_HIDDEN, hs),
+            )
+
+        def forward(self, input_ids, attention_mask):
+            out = self.base(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+                return_dict=True
+            )
+            hidden = out.hidden_states[-1]        # (B, T, H)
+            emb    = hidden.mean(dim=1)           # mean-pooling
+            emb    = torch.nan_to_num(emb, nan=0.0, posinf=1e-6, neginf=-1e-6)
+            z      = self.proj(emb)               # now (B, H)
+            z      = torch.nan_to_num(z, nan=0.0, posinf=1e-6, neginf=-1e-6)
+            norm   = z.norm(p=2, dim=1, keepdim=True).clamp_min(1e-6)
+            return z / norm
+
+    model = GemmaSimCSE(model_lora).to(device)
+    torch.autograd.set_detect_anomaly(True)
+
+    # Load and split dataset
+    raw = load_dataset("text", data_files={"train": DATA_FILE}, split="train")
+    raw = raw.filter(lambda x: x["text"].strip() != "")
+    split = raw.train_test_split(test_size=VAL_RATIO, seed=SEED)
+    train_ds = split["train"]
+    val_ds   = split["test"]
+
+    # Tokenization
+    def tokenize_fn(batch):
+        toks = tokenizer(
+            batch["text"],
+            max_length=MAX_LENGTH,
+            truncation=True,
+            padding="max_length"
+        )
+        return {"input_ids": toks["input_ids"], "attention_mask": toks["attention_mask"]}
+
+    train_ds = train_ds.map(
+        tokenize_fn,
+        batched=True,
+        batch_size=1000,
+        num_proc=4,
+        remove_columns=["text"]
+    )
+    val_ds = val_ds.map(
+        tokenize_fn,
+        batched=True,
+        batch_size=1000,
+        num_proc=4,
+        remove_columns=["text"]
+    )
+
+    train_ds.set_format(type="torch", columns=["input_ids", "attention_mask"])
+    val_ds.set_format(type="torch", columns=["input_ids", "attention_mask"])
+
+    train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)
+    val_loader   = DataLoader(val_ds,   batch_size=BATCH_SIZE, shuffle=False)
+
+    # Optimizer & scheduler
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY
+    )
+    total_steps = len(train_loader) * NUM_EPOCHS
+    scheduler = get_linear_schedule_with_warmup(
+        optimizer,
+        num_warmup_steps=int(0.1 * total_steps),
+        num_training_steps=total_steps
+    )
+
+    # Training + validation loop
+    for epoch in range(1, NUM_EPOCHS + 1):
+        # --- train ---
+        model.train()
+        train_loss = 0.0
+        for batch in tqdm(train_loader, desc=f"Train Epoch {epoch}", unit="batch"):
+            ids  = batch["input_ids"].to(device)
+            mask = batch["attention_mask"].to(device)
+
+            emb1 = model(ids, mask)
+            emb2 = model(ids, mask)
+            emb  = torch.cat([emb1, emb2], dim=0)
+            sim  = (emb @ emb.T) / TEMP
+            sim  = sim.clamp(SIM_CLAMP_MIN, SIM_CLAMP_MAX)
+            sim.fill_diagonal_(-1e9)
+
+            B = emb1.size(0)
+            labels = torch.cat([
+                torch.arange(B, device=device) + B,
+                torch.arange(B, device=device)
+            ], dim=0)
+
+            loss = F.cross_entropy(sim, labels)
+            optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP_NORM)
+            optimizer.step()
+            scheduler.step()
+            train_loss += loss.item()
+
+        avg_train_loss = train_loss / len(train_loader)
+        print(f"Epoch {epoch} training complete. avg train loss: {avg_train_loss:.6f}")
+
+        # --- validate ---
+        model.eval()
+        val_loss = 0.0
+        with torch.no_grad():
+            for batch in tqdm(val_loader, desc=f"Validate Epoch {epoch}", unit="batch"):
+                ids  = batch["input_ids"].to(device)
+                mask = batch["attention_mask"].to(device)
+
+                emb1 = model(ids, mask)
+                emb2 = model(ids, mask)
+                emb  = torch.cat([emb1, emb2], dim=0)
+                sim  = (emb @ emb.T) / TEMP
+                sim  = sim.clamp(SIM_CLAMP_MIN, SIM_CLAMP_MAX)
+                sim.fill_diagonal_(-1e9)
+
+                B = emb1.size(0)
+                labels = torch.cat([
+                    torch.arange(B, device=device) + B,
+                    torch.arange(B, device=device)
+                ], dim=0)
+
+                loss = F.cross_entropy(sim, labels)
+                val_loss += loss.item()
+
+        avg_val_loss = val_loss / len(val_loader)
+        print(f"Epoch {epoch} validation complete. avg val loss: {avg_val_loss:.6f}")
+
+        # save checkpoint
+        ckpt_dir = os.path.join(OUTPUT_DIR, f"epoch{epoch}")
+        model_lora.save_pretrained(ckpt_dir)
+        tokenizer.save_pretrained(ckpt_dir)
+
+    # save final model
+    final_dir = os.path.join(OUTPUT_DIR, "final")
+    os.makedirs(final_dir, exist_ok=True)
+    model_lora.save_pretrained(final_dir)
+    tokenizer.save_pretrained(final_dir)
+    print("Training and validation complete. Final model saved to", final_dir)
+
+if __name__ == "__main__":
+    freeze_support()
+    main()