Skip to content

Revert "Feature/reconstruction accuracy" #10

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
knc6 merged 1 commit into from
Oct 10, 2025
Merged

Revert "Feature/reconstruction accuracy" #10

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
227 changes: 53 additions & 174 deletions atomgpt/inverse_models/inverse_models.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
)
from transformers import (
TrainingArguments,
EarlyStoppingCallback,
)
import torch
from atomgpt.inverse_models.utils import (
Expand All @@ -18,7 +17,6 @@
get_figlet,
)
from trl import SFTTrainer, SFTConfig
from trl import DataCollatorForCompletionOnlyLM # ADDED
from peft import PeftModel
from datasets import load_dataset
from functools import partial
Expand All @@ -29,7 +27,6 @@
from jarvis.io.vasp.inputs import Poscar
import csv
import os
import numpy as np
from pydantic_settings import BaseSettings
import sys
import json
Expand Down Expand Up @@ -67,14 +64,14 @@ class TrainingPropConfig(BaseSettings):
gradient_accumulation_steps: int = 4
num_train: Optional[int] = None
num_test: Optional[int] = None
test_ratio: Optional[float] = 0.2
test_ratio: Optional[float] = 0.1
val_ratio: Optional[float] = 0.1
model_save_path: str = "atomgpt_lora_model"
lora_rank: Optional[int] = 16
lora_alpha: Optional[int] = 16
loss_type: str = "default"
optim: str = "adamw_8bit"
id_tag: str = "id"
save_strategy: str = "st"
lr_scheduler_type: str = "linear"
separator: str = ","
prop: str = "Tc_supercon"
Expand All @@ -101,16 +98,15 @@ class TrainingPropConfig(BaseSettings):
output_prompt: str = (
" Generate atomic structure description with lattice lengths, angles, coordinates and atom types."
)
val_ratio: Optional[float] = 0.1
num_val: Optional[int] = None
eval_strategy: Literal["no", "steps", "epoch"] = "steps"
eval_steps: int = 200
load_best_model_at_end: bool = True
metric_for_best_model: str = "eval_loss"
greater_is_better: bool = False
save_total_limit: int = 2
early_stopping_patience: int = 3
early_stopping_threshold: float = 0.0
# num_val: Optional[int] = 2
hp_cfg_path: Optional[str] = "hp_search_config.json"
per_device_train_batch_size: int = 2
gradient_accumulation_steps: int = 4
warmup_steps: int = 3
warmup_ratio: float = 0.0
logging_steps: int = 10




def get_input(config=None, chem="", val=10):
Expand Down Expand Up @@ -217,96 +213,42 @@ def load_model(path="", config=None):
FastLanguageModel.for_inference(model)
return model, tokenizer, config

def _validate_atoms(atoms):
if atoms is None:
return False, "atoms_is_none"
try:
lat = np.asarray(getattr(atoms, "lattice_mat", None), dtype=float)
if lat.shape != (3, 3):
return False, f"bad_lattice_shape:{getattr(atoms,'lattice_mat',None)}"
if not np.isfinite(lat).all():
return False, "nonfinite_lattice"
n = getattr(atoms, "num_atoms", None)
if n is None or n <= 0:
return False, f"num_atoms_invalid:{n}"
_ = Poscar(atoms).to_string()
return True, ""
except Exception as e:
return False, f"poscar_fail:{type(e).__name__}:{e}"

def _poscar_one_line(at):
return Poscar(at).to_string().replace("\n", "\\n")

def _misses_path(csv_out, config):
fname = getattr(config, "miss_csv", None)
if fname is None or not str(fname).strip():
root, ext = os.path.splitext(csv_out)
fname = root + ".misses.csv"
os.makedirs(os.path.dirname(os.path.abspath(fname)), exist_ok=True)
return fname

def evaluate(
test_set=[],
model="",
tokenizer="",
csv_out="out.csv",
config="",
test_set=[], model="", tokenizer="", csv_out="out.csv", config=""
):
print("Testing\n", len(test_set))
os.makedirs(os.path.dirname(os.path.abspath(csv_out)), exist_ok=True)
miss_csv_out = _misses_path(csv_out, config)

with open(csv_out, "w", newline="") as f_ok, open(miss_csv_out, "w", newline="") as f_miss:
ok_writer = csv.writer(f_ok)
miss_writer = csv.writer(f_miss)
ok_writer.writerow(["id", "target", "prediction"])
miss_writer.writerow(["id", "stage", "error", "detail", "raw_text_preview"])

for i in tqdm(test_set, total=len(test_set)):
sample_id = i.get("id", "")
target_mat = None
target_err = None
try:
target_mat = text2atoms("\n" + i["output"])
ok, detail = _validate_atoms(target_mat)
if not ok:
target_err = detail
except Exception as e:
target_err = f"text2atoms:{type(e).__name__}:{e}"

if target_err:
miss_writer.writerow([sample_id, "target", "invalid_target", target_err, (i.get("output","")[:240])])
continue

gen_mat = None
gen_err = None
try:
gen_mat = gen_atoms(
prompt=i["input"],
tokenizer=tokenizer,
model=model,
alpaca_prompt=config.alpaca_prompt,
instruction=config.instruction,
)
ok, detail = _validate_atoms(gen_mat)
if not ok:
gen_err = detail
except Exception as e:
gen_err = f"gen_atoms:{type(e).__name__}:{e}"

if gen_err:
miss_writer.writerow([sample_id, "prediction", "invalid_prediction", gen_err, ""])
continue

try:
ok_writer.writerow([
sample_id,
_poscar_one_line(target_mat),
_poscar_one_line(gen_mat),
])
except Exception as e:
miss_writer.writerow([sample_id, "write", "write_failed", f"{type(e).__name__}:{e}", ""])
f = open(csv_out, "w")
f.write("id,target,prediction\n")

for i in tqdm(test_set, total=len(test_set)):
# try:
# prompt = i["input"]
# print("prompt", prompt)
gen_mat = gen_atoms(
prompt=i["input"],
tokenizer=tokenizer,
model=model,
alpaca_prompt=config.alpaca_prompt,
instruction=config.instruction,
)
target_mat = text2atoms("\n" + i["output"])
print("target_mat", target_mat)
print("genmat", gen_mat)
line = (
i["id"]
+ ","
+ Poscar(target_mat).to_string().replace("\n", "\\n")
+ ","
+ Poscar(gen_mat).to_string().replace("\n", "\\n")
+ "\n"
)
f.write(line)
# print()
# except Exception as exp:
# print("Error", exp)
# pass
f.close()


def batch_evaluate(
Expand Down Expand Up @@ -422,7 +364,6 @@ def main(config_file=None):
run_path = os.path.dirname(id_prop_path)
num_train = config.num_train
num_test = config.num_test
num_val = config.num_val
# model_name = config.model_name
callback_samples = config.callback_samples
# loss_function = config.loss_function
Expand All @@ -431,12 +372,8 @@ def main(config_file=None):
reader = csv.reader(f)
dt = [row for row in reader]
if not num_train:
if num_test is None:
num_test = int(len(dt) * config.test_ratio)
if num_val is None:
vr = config.val_ratio if config.val_ratio is not None else 0.1
num_val = int(len(dt) * vr)
num_train = len(dt) - num_test - num_val
num_test = int(len(dt) * config.test_ratio)
num_train = len(dt) - num_test

dat = []
ids = []
Expand All @@ -454,7 +391,7 @@ def main(config_file=None):
# if ";" in i[1]:
# tmp = "\n".join([str(round(float(j), 2)) for j in i[1].split(";")])
# else:
# tmp = str(round(float(i[1]), 3))
# tmp = str(round(float(i[1]), 3))
info[config.prop] = (
tmp # float(i[1]) # [float(j) for j in i[1:]] # float(i[1]
)
Expand All @@ -473,10 +410,8 @@ def main(config_file=None):

train_ids = ids[0:num_train]
print("num_train", num_train)
print("num_val", num_val)
print("num_test", num_test)
val_ids = ids[num_train : num_train + num_val]
test_ids = ids[num_train + num_val : num_train + num_val + num_test]
test_ids = ids[num_train : num_train + num_test]
# test_ids = ids[num_train:]
alpaca_prop_train_filename = os.path.join(
config.output_dir, "alpaca_prop_train.json"
Expand All @@ -497,20 +432,6 @@ def main(config_file=None):
m_train = loadjson(alpaca_prop_train_filename)
print("Sample:\n", m_train[0])

alpaca_prop_val_filename = os.path.join(
config.output_dir, "alpaca_prop_val.json"
)
if not os.path.exists(alpaca_prop_val_filename):
m_val = make_alpaca_json(
dataset=dat,
jids=val_ids,
config=config,
)
dumpjson(data=m_val, filename=alpaca_prop_val_filename)
else:
print(alpaca_prop_val_filename, "exists")
m_val = loadjson(alpaca_prop_val_filename)

alpaca_prop_test_filename = os.path.join(
config.output_dir, "alpaca_prop_test.json"
)
Expand Down Expand Up @@ -573,11 +494,6 @@ def main(config_file=None):
split="train",
# "json", data_files="alpaca_prop_train.json", split="train"
)
val_dataset = load_dataset(
"json",
data_files=alpaca_prop_val_filename,
split="train",
)
eval_dataset = load_dataset(
"json",
data_files=alpaca_prop_test_filename,
Expand All @@ -600,10 +516,6 @@ def tokenize_function(example):
formatting_prompts_func_with_prompt,
batched=True,
)
val_dataset = val_dataset.map(
formatting_prompts_func_with_prompt,
batched=True,
)
eval_dataset = eval_dataset.map(
formatting_prompts_func_with_prompt,
batched=True,
Expand All @@ -617,24 +529,14 @@ def tokenize_function(example):
print(f"🧠 Suggested max_seq_length based on dataset: {max_seq_length}")

tokenized_train = train_dataset.map(tokenize_function, batched=True)
tokenized_val = val_dataset.map(tokenize_function, batched=True)
tokenized_eval = eval_dataset.map(tokenize_function, batched=True)
tokenized_train.set_format(
type="torch", columns=["input_ids", "attention_mask", "output"]
)
tokenized_val.set_format(
type="torch", columns=["input_ids", "attention_mask", "output"]
)
tokenized_eval.set_format(
type="torch", columns=["input_ids", "attention_mask", "output"]
)

# ADDED: Completion-only collator to mask loss to the completion after "### Output:"
response_template = "### Output:" # ADDED
collator = DataCollatorForCompletionOnlyLM( # ADDED
response_template=response_template, tokenizer=tokenizer, mlm=False
) # ADDED

"""
trainer = SFTTrainer(
# trainer = CustomSFTTrainer(
Expand Down Expand Up @@ -673,42 +575,30 @@ def tokenize_function(example):

trainer = SFTTrainer(
model=model,
train_dataset=train_dataset, # ADDED/MODIFIED: use raw text dataset so collator can mask
eval_dataset=val_dataset, # ADDED/MODIFIED
processing_class=tokenizer, # ADDED
data_collator=collator, # ADDED
train_dataset=tokenized_train,
# train_dataset = train_dataset,
# tokenizer = tokenizer,
args=SFTConfig(
dataset_text_field="text",
max_seq_length=config.max_seq_length,
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
warmup_steps=5,
per_device_train_batch_size=config.per_device_train_batch_size,
gradient_accumulation_steps=config.gradient_accumulation_steps,
warmup_steps=config.warmup_steps,
overwrite_output_dir=True,
warmup_ratio=config.warmup_ratio,
# max_steps=60,
logging_steps=30,
logging_steps=config.logging_steps,
output_dir=config.output_dir,
optim=config.optim,
seed=config.seed_val,
num_train_epochs=config.num_epochs,
save_strategy=config.save_strategy,
save_steps=config.save_steps,
eval_strategy=config.eval_strategy,
eval_steps=config.eval_steps,
load_best_model_at_end=config.load_best_model_at_end,
metric_for_best_model=config.metric_for_best_model,
greater_is_better=config.greater_is_better,
save_total_limit=config.save_total_limit,
dataloader_num_workers=8,
dataloader_pin_memory=True,
dataloader_persistent_workers=True,
group_by_length=True,
),
)
if callback_samples > 0:
callback = ExampleTrainerCallback(
some_tokenized_dataset=tokenized_val, # use validation set
some_tokenized_dataset=tokenized_eval,
# some_tokenized_dataset=tokenized_eval,
tokenizer=tokenizer,
max_length=config.max_seq_length,
Expand All @@ -717,10 +607,6 @@ def tokenize_function(example):
trainer.add_callback(callback)
gpu_usage = PrintGPUUsageCallback()
trainer.add_callback(gpu_usage)
trainer.add_callback(EarlyStoppingCallback(
early_stopping_patience=config.early_stopping_patience,
early_stopping_threshold=config.early_stopping_threshold,
))
trainer_stats = trainer.train()
trainer.save_model(config.model_save_path)
# model.save_pretrained(config.model_save_path)
Expand All @@ -733,13 +619,6 @@ def tokenize_function(example):
# )
model = trainer.model
FastLanguageModel.for_inference(model) # Enable native 2x faster inference

# ADDED: Set greedy decoding defaults with explicit stops (affects gen_atoms & generate)
model.generation_config.do_sample = False
model.generation_config.eos_token_id = tokenizer.eos_token_id
model.generation_config.pad_token_id = tokenizer.eos_token_id
model.generation_config.max_new_tokens = min(512, config.max_seq_length)

# model, tokenizer, config = load_model(path=config.model_save_path)
# batch_evaluate(
# prompts=[i["input"] for i in m_test],
Expand Down
Loading