[Python] Trying to train/develop model for email categorizing, not getting the desired result

Heyy! I was trying to train a model for email categorizing, it only needs to give me a folder path, which i will put in the newly arrived email using a script.

It's written in hungarian (my first language) because i want to use it for hungarian emails.

i've used around 20-30k emails to train the model. During the training, the loss went just under 1 by the end of the training, but the results were not really good. It was giving me either the wrong folder paths or not giving any outputs.

(I've never trained a model in my life before this. I'm not really familiar with training ai models. Also i haven't really used stackoverflow before, so i'm really sorry if i posted in a wrong format, or not using correct language)

Model i used for the training: https://huggingface.co/Cseti/PULI-LlumiX-Llama-3.1-Instruct-LoRA

i used the following code for the training


from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
TrainingArguments,
Trainer,
BitsAndBytesConfig,
DataCollatorForLanguageModeling,
TrainerCallback # <--- Ez a hiányzó import
)
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from datasets import load_dataset
import torch
import os

# 1. Memóriaoptimalizált modellbetöltés
model_name = "Cseti/PULI-LlumiX-Llama-3.1-Instruct-LoRA"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token

bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True
)

# Modell betöltése és előkészítése
model = AutoModelForCausalLM.from_pretrained(
model_name,
quantization_config=bnb_config,
device_map="auto",
use_cache=False,
torch_dtype=torch.float16,
)
model = prepare_model_for_kbit_training(model)

# 2. Optimalizált LoRA konfiguráció
peft_config = LoraConfig(
r=32,
lora_alpha=64,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
lora_dropout=0.1,
bias="none",
task_type="CAUSAL_LM"
)
model = get_peft_model(model, peft_config)

# 3. Dinamikus adatfeldolgozás (MÓDOSÍTVA)
class ProgressCallback(TrainerCallback):
def on_log(self, args, state, control, logs=None, **kwargs):
if state.epoch is not None:
# Tanítási loss
if "loss" in logs:
print(f"[TRAIN] Epoch: {state.epoch+1} | Step: {state.global_step} | Loss: {logs['loss']:.4f}")
# Validációs loss
if "eval_loss" in logs:
print(f"[EVAL] Epoch: {state.epoch+1} | Step: {state.global_step} | Loss: {logs['eval_loss']:.4f}")

data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer,
mlm=False,
pad_to_multiple_of=8
)

def process_dataset(examples):
batch = []
for msg in examples["messages"]:
text = tokenizer.apply_chat_template(msg, tokenize=False)
batch.append(text)
return tokenizer(
batch,
max_length=2048,
truncation=True,
padding="max_length",
)

# Adatok betöltése és felosztása
dataset = load_dataset("json", data_files="finomhangolasi_adatok_2023_szept.jsonl", split="train")
dataset = dataset.train_test_split(test_size=0.1)

# 4. Optimalizált tanítási paraméterek
training_args = TrainingArguments(
output_dir="./finomhangolt_modell",
per_device_train_batch_size=4,
gradient_accumulation_steps=8,
learning_rate=3e-4,
num_train_epochs=5,
weight_decay=0.01,
bf16=True,
gradient_checkpointing=True,
optim="adamw_torch",
logging_steps=20,
save_strategy="steps",
save_steps=200,
evaluation_strategy="steps",
eval_steps=200,
report_to="none",
dataloader_num_workers=8,
max_grad_norm=0.3,
warmup_ratio=0.1
)

# 5. Trainer konfiguráció
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"],
data_collator=data_collator,
callbacks=[ProgressCallback()] # Módosított callback
)

# 6. Tanítás indítása
print(" Tanítás indítása...")
trainer.train()

# 7. Memóriaoptimalizált mentés
model.save_pretrained(
"./finomhangolt_modell_teszt2",
safe_serialization=True,
max_shard_size="2GB"
)
print("✅ Modell sikeresen elmentve!")


I expected it to work. I've tested it with some other emails that i excluded from the training jsonl file, but it was getting the folder path all wrong.

This is the testing code i tried to run:

from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
BitsAndBytesConfig
)
from peft import PeftModel
import torch

BASE_MODEL = "Cseti/PULI-LlumiX-Llama-3.1-Instruct-LoRA"
ADAPTER_PATH = "./finomhangolt_modell_teszt2"

bnb_config = BitsAndBytesConfig(
load_in_4bit=True,
bnb_4bit_quant_type="nf4",
bnb_4bit_compute_dtype=torch.float16,
bnb_4bit_use_double_quant=True
)

tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
tokenizer.pad_token = tokenizer.eos_token

model = AutoModelForCausalLM.from_pretrained(
BASE_MODEL,
quantization_config=bnb_config,
device_map="auto",
torch_dtype=torch.float16,
)
model = PeftModel.from_pretrained(model, ADAPTER_PATH)
model.eval()

def build_chat_prompt(email_text):
return [
{"role": "system", "content": "Te egy email kategorizáló AI vagy. Csak a mappa elérési útját add vissza."},
{"role": "user", "content": email_text}
]

email_text = """Tárgy: .................
Tartalom: .................................
"""

prompt = build_chat_prompt(email_text)
prompt_text = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True)

inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
with torch.no_grad():
outputs = model.generate(
**inputs,
max_new_tokens=20,
do_sample=False,
pad_token_id=tokenizer.eos_token_id,
early_stopping=True
)

full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
response = full_response[len(prompt_text):].strip().split("\n")[0]

print(" Modell által javasolt mappa:")
print(response)



I had to blur out a part where the email is because of privacy reasons. I really have no clue if it's the training i did was not good or the test code i ran was poor.

Continue reading...