From 74b2cc519765dc9aa42134c08e23782d3d0176b2 Mon Sep 17 00:00:00 2001 From: Kamal Choudhary Date: Fri, 10 Oct 2025 11:50:16 -0400 Subject: [PATCH] Revert "Feature/reconstruction accuracy" --- atomgpt/inverse_models/inverse_models.py | 227 ++++++----------------- 1 file changed, 53 insertions(+), 174 deletions(-) diff --git a/atomgpt/inverse_models/inverse_models.py b/atomgpt/inverse_models/inverse_models.py index 89cdad8..7778ec3 100644 --- a/atomgpt/inverse_models/inverse_models.py +++ b/atomgpt/inverse_models/inverse_models.py @@ -8,7 +8,6 @@ ) from transformers import ( TrainingArguments, - EarlyStoppingCallback, ) import torch from atomgpt.inverse_models.utils import ( @@ -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 @@ -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 @@ -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" @@ -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): @@ -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( @@ -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 @@ -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 = [] @@ -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] ) @@ -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" @@ -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" ) @@ -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, @@ -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, @@ -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( @@ -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, @@ -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) @@ -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],