Update files

Changed base model to Gwen 2.5 Coder
Changed epochs to 5.
Changed final model folder name.
Updated READMEs

README.es.md

@@ -11,7 +11,7 @@ tags:
 license: "apache-2.0"
 datasets:
 - xlangai/spider
-base_model: "deepseek-ai/deepseek-coder-1.3b-base"
+base_model: "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 ---
 
 > [Ver en ingles/See in english](https://huggingface.co/Komma-LuisMiSanVe/LangToSQL/blob/main/README.md)
@@ -38,24 +38,27 @@ base_model: "deepseek-ai/deepseek-coder-1.3b-base"
 El modelo de IA ha sido entrenado para convertir lenguaje natural a sentencias de PostgreSQL.
 
 ## 📝 Explicación de Tecnología
-El modelo usa [DeepSeek Coder](https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-base) de base y refinado con los datasets de [Spider](https://yale-lily.github.io/spider).
+El modelo usa [Gwen Coder](https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct) de base y refinado con los datasets de [Spider](https://yale-lily.github.io/spider).
 
 El dataset en archivo `JSON` contiene `train_spider.json` de **Spider**, ya que es el dataset principal.
 
-El modelo se puede exportar a `GGUF` con [llama.cpp](https://github.com/ggml-org/llama.cpp) para que puedas usarlo en programas como [LM Studio](https://lmstudio.ai/).
+El modelo se ha exportado a `GGUF` con [llama.cpp](https://github.com/ggml-org/llama.cpp) para que puedas usarlo en programas como [LM Studio](https://lmstudio.ai/).
 
 ## 🛠️ Instalación
 Para ejecutar el script de entrenamiento por tu cuenta, primero necesitas instalar [Python](https://www.python.org/) y ejecuta este comando:
 ```
-pip install transformers datasets peft accelerate bitsandbytes trl
+pip install transformers datasets peft accelerate bitsandbytes trl==1.0.0
 ```
 Dependiendo en la versión, es posible que necesites usar este en su lugar:
 ```
-py -m pip install transformers datasets peft accelerate bitsandbytes trl
+py -m pip install transformers datasets peft accelerate bitsandbytes trl==1.0.0
 ```
 
+>[!IMPORTANT]
+>Asegurate que la libreria `TRL` está en la versión `1.0.0`, ya que es la única version compatible con el script de entrenamiento.
+
 ## 📂 Archivos
-Este repositorio incluye los archivos del modelo LLM entrenado, su script de entrenamiento y el dataset para entrenar.
+Este repositorio incluye los archivos del modelo LLM entrenado, su script de entrenamiento, el dataset para entrenar y un script para probar el modelo `.safetensors`.
 
 Puedes descargar el `GGUF` final desde los [Lanzamientos](https://github.com/LuisMiSanVe/LangToSQL_LLM/releases).
 
@@ -79,6 +82,6 @@ El número de la versión seguirá este formato: \
   - [trl](https://pypi.org/project/trl/)
 - Otros:
   - [llama.cpp](https://lmstudio.ai/)
-  - [DeepSeek Coder](https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-base)
+  - [Gwen Coder](https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct)
   - [Spider](https://yale-lily.github.io/spider)
 - IDE Recomendado: [VS Code](https://code.visualstudio.com/)

README.md

@@ -11,7 +11,7 @@ tags:
 license: "apache-2.0"
 datasets:
 - xlangai/spider
-base_model: "deepseek-ai/deepseek-coder-1.3b-base"
+base_model: "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 ---
 
 > [See in spanish/Ver en español](https://huggingface.co/Komma-LuisMiSanVe/LangToSQL/blob/main/README.es.md)
@@ -38,24 +38,27 @@ base_model: "deepseek-ai/deepseek-coder-1.3b-base"
 The AI model has been trained for turning natural language to PostgreSQL queries.
 
 ## 📝 Technology Explanation
-This model uses [DeepSeek Coder](https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-base) as a base and then is fine tuned with [Spider](https://yale-lily.github.io/spider) datasets.
+This model uses [Gwen Coder](https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct) as a base and then is fine tuned with [Spider](https://yale-lily.github.io/spider) datasets.
 
 The `JSON` dataset file contains **Spider**'s `train_spider.json` as is the main dataset.
 
-The model can be exported to `GGUF` with [llama.cpp](https://github.com/ggml-org/llama.cpp) so it can be used by programs like [LM Studio](https://lmstudio.ai/).
+The model is exported to `GGUF` with [llama.cpp](https://github.com/ggml-org/llama.cpp) so it can be used by programs like [LM Studio](https://lmstudio.ai/).
 
 ## 🛠️ Setup
 In order to execute the training script for your own, you first need to install [Python](https://www.python.org/) and run this command:
 ```
-pip install transformers datasets peft accelerate bitsandbytes trl
+pip install transformers datasets peft accelerate bitsandbytes trl==1.0.0
 ```
 Depending on the version, you may have to use this instead:
 ```
-py -m pip install transformers datasets peft accelerate bitsandbytes trl
+py -m pip install transformers datasets peft accelerate bitsandbytes trl==1.0.0
 ```
 
+>[!IMPORTANT]
+>Make sure the `TRL` library version is `1.0.0`, as is the only version supported by the trainer script.
+
 ## 📂 Files
-This repository includes the trained LLM model's files, its training script and the training dataset.
+This repository includes the trained LLM model's files, its training script, the training dataset and a tester script to test the `.safetensors` model.
 
 You can download the final `GGUF` in the [Releases](https://github.com/LuisMiSanVe/LangToSQL_LLM/releases).
 
@@ -76,9 +79,9 @@ The version number will follow this format: \
   - [peft](https://pypi.org/project/peft/)
   - [acceletare](https://pypi.org/project/accelerate/)
   - [bitsandbytes](https://pypi.org/project/bitsandbytes/)
-  - [trl](https://pypi.org/project/trl/)
+  - [trl](https://pypi.org/project/trl/) (1.0.0)
 - Other: 
   - [llama.cpp](https://github.com/ggml-org/llama.cpp)
-  - [DeepSeek Coder](https://huggingface.co/deepseek-ai/deepseek-coder-1.3b-base)
+  - [Gwen Coder](https://huggingface.co/Qwen/Qwen2.5-Coder-1.5B-Instruct)
   - [Spider](https://yale-lily.github.io/spider)
 - Recommended IDE: [VS Code](https://code.visualstudio.com/)

test.py

@@ -0,0 +1,49 @@
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+MODEL_PATH = "./sql-model-merged"  
+
+PROMPT = """\
+Write a select query of the invoice table.
+"""
+
+print("Loading tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_PATH
+)
+
+print("Loading model... (this may take a while)")
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_PATH,
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    device_map="auto",
+    ignore_mismatched_sizes=True
+)
+
+model.eval()
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+inputs = tokenizer(PROMPT, return_tensors="pt").to(model.device)
+
+print("\nGenerating response...\n")
+
+with torch.no_grad():
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=256,
+        temperature=0.2,
+        top_p=0.95,
+        do_sample=True,
+        repetition_penalty=1.1,
+        eos_token_id=tokenizer.eos_token_id
+    )
+
+result = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+print("===== MODEL OUTPUT =====\n")
+print(result)
+print("\n========================")

trainer.py

@@ -4,52 +4,35 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments
 from peft import LoraConfig, PeftModel
 from trl import SFTTrainer
 
-model_name = "deepseek-ai/deepseek-coder-1.3b-base"
+model_name = "Qwen/Qwen2.5-Coder-1.5B-Instruct"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 tokenizer.pad_token = tokenizer.eos_token
 
 model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float32,
-    device_map={"": "cpu"} # Sets CPU for training, you can change it to use the GPU instead
+    device_map="auto"
 )
 
+model.config.pad_token_id = tokenizer.eos_token_id
+
 dataset = load_dataset("json", data_files="train.json", split="train")
 
-def format_example(example):
+def format_example(x):
+    messages = [
+        {"role": "user", "content": f"Write SQL query for: {x['question']}"},
+        {"role": "assistant", "content": x["query"]}
+    ]
     return {
-        "instruction": example["question"],
-        "input": "",
-        "output": example["query"]
+        "text": tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False
+        )
     }
 
 dataset = dataset.map(format_example)
 
-def tokenize(example):
-    prompt_ids = tokenizer(
-        example["instruction"],
-        padding="max_length",
-        truncation=True,
-        max_length=512
-    ).input_ids
-
-    label_ids = tokenizer(
-        example["output"],
-        padding="max_length",
-        truncation=True,
-        max_length=512
-    ).input_ids
-
-    attention_mask = [1 if id != tokenizer.pad_token_id else 0 for id in prompt_ids]
-
-    return {
-        "input_ids": prompt_ids,
-        "attention_mask": attention_mask,
-        "labels": label_ids
-    }
-
-dataset = dataset.map(tokenize, batched=False)
-
 peft_config = LoraConfig(
     r=16,
     lora_alpha=32,
@@ -64,10 +47,10 @@ training_args = TrainingArguments(
     per_device_train_batch_size=1,
     gradient_accumulation_steps=4,
     learning_rate=2e-4,
-    num_train_epochs=1, # More epochs -> better accuracy but longer training
+    num_train_epochs=5,
     logging_steps=10,
     save_strategy="epoch",
-    fp16=False
+    fp16=torch.cuda.is_available()
 )
 
 trainer = SFTTrainer(
@@ -85,8 +68,9 @@ tokenizer.save_pretrained("./sql-model")
 base_model = AutoModelForCausalLM.from_pretrained(
     model_name,
     torch_dtype=torch.float32,
-    device_map={"": "cpu"}
+    device_map="auto"
 )
-model_merged = PeftModel.from_pretrained(base_model, "./sql-model")
-model_merged = model_merged.merge_and_unload()
-model_merged.save_pretrained("./sql-model-merged")
+model = PeftModel.from_pretrained(base_model, "./sql-model")
+model = model.merge_and_unload()
+model.save_pretrained("./sql-model-merged", safe_serialization=True)
+tokenizer.save_pretrained("./sql-model-merged")