diff --git a/atomgpt/inverse_models/inverse_models.py b/atomgpt/inverse_models/inverse_models.py index 9744286..8388265 100644 --- a/atomgpt/inverse_models/inverse_models.py +++ b/atomgpt/inverse_models/inverse_models.py @@ -27,6 +27,7 @@ from jarvis.io.vasp.inputs import Poscar import csv import os +import numpy as np from pydantic_settings import BaseSettings import sys import json @@ -211,42 +212,96 @@ 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)) - f = open(csv_out, "w") - f.write("id,target,prediction\n") + 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}", ""]) - 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( @@ -514,10 +569,12 @@ def tokenize_function(example): train_dataset = train_dataset.map( formatting_prompts_func_with_prompt, batched=True, + num_proc=config.dataset_num_proc ) eval_dataset = eval_dataset.map( formatting_prompts_func_with_prompt, batched=True, + num_proc=config.dataset_num_proc ) # Compute the actual max sequence length in raw text lengths = [ @@ -527,8 +584,8 @@ def tokenize_function(example): max_seq_length = max(lengths) print(f"🧠 Suggested max_seq_length based on dataset: {max_seq_length}") - tokenized_train = train_dataset.map(tokenize_function, batched=True) - tokenized_eval = eval_dataset.map(tokenize_function, batched=True) + tokenized_train = train_dataset.map(tokenize_function, batched=True, num_proc=config.dataset_num_proc) + tokenized_eval = eval_dataset.map(tokenize_function, batched=True, num_proc=config.dataset_num_proc) tokenized_train.set_format( type="torch", columns=["input_ids", "attention_mask", "output"] ) diff --git a/atomgpt/inverse_models/inverse_predict.py b/atomgpt/inverse_models/inverse_predict.py index 378ca35..63054cf 100644 --- a/atomgpt/inverse_models/inverse_predict.py +++ b/atomgpt/inverse_models/inverse_predict.py @@ -79,13 +79,10 @@ def relax_atoms( calculator = AlignnAtomwiseCalculator(path=default_path(), device="cpu") t1 = time.time() - # if calculator is None: - # return atoms ase_atoms = atoms.ase_converter() ase_atoms.calc = calculator ase_atoms = ExpCellFilter(ase_atoms, constant_volume=constant_volume) - # TODO: Make it work with any other optimizer dyn = FIRE(ase_atoms) dyn.run(fmax=fmax, steps=nsteps) en = ase_atoms.atoms.get_potential_energy() @@ -114,9 +111,6 @@ def predict( load_in_4bit=None, # temp_config["load_in_4bit"] verbose=True, # temp_config["load_in_4bit"] ): - # if not os.path.exists("config_name"): - - # config_name=os.path.join(output_dir,"config.json") print("config_path", config_path) if output_dir is not None: config_name = os.path.join(output_dir, "config.json") @@ -125,7 +119,7 @@ def predict( config_name = os.path.join(parent, "config.json") adapter = os.path.join(output_dir, "adapter_config.json") if os.path.exists(adapter): - model_name = output_dir # temp_config["model_name"] + model_name = output_dir if config_path is not None: config_name = config_path if verbose: @@ -198,21 +192,18 @@ def predict( formula=formula, background_subs=background_subs, ) - # y[y < 0.1] = 0 - y_new_str = y # "\n".join(["{0:.2f}".format(x) for x in y]) + y_new_str = y try: if ".dat" in i: formula = str(_formula.split("/")[-1].split(".dat")[0]) except Exception: pass - # gen_mat = main_spectra(spectra=[[y_new_str,y]],formulas=[formula],model=model,tokenizer=tokenizer,device='cuda')[0] prompt = ( "The chemical formula is " + formula + " The " + temp_config["prop"] + " is " - # + " The XRD is " + y_new_str + ". Generate atomic structure description with lattice lengths, angles, coordinates and atom types." ) @@ -224,13 +215,13 @@ def predict( + " The " + temp_config["prop"] + " is " - # + " The XRD is " + str(prop_val) + ". Generate atomic structure description with lattice lengths, angles, coordinates and atom types." ) if verbose: print("prompt here", prompt.replace("\n", ",")) + gen_mat = gen_atoms( prompt=prompt, model=model, @@ -239,6 +230,18 @@ def predict( instruction=temp_config["instruction"], device=device, ) + + if gen_mat is None: + print( + "The structure returned by gen_mat() is not a valid crystal structure." + ) + info = {} + info["prompt"] = prompt + info["error"] = "Invalid structure returned by AtomGPT (None)." + mem.append(info) + # skip the rest of the loop for this entry + continue + if verbose: print("gen atoms", gen_mat) print("gen atoms spacegroup", gen_mat.spacegroup()) @@ -257,8 +260,6 @@ def predict( if __name__ == "__main__": - # output_dir = make_id_prop() - # output_dir="." args = parser.parse_args(sys.argv[1:]) print("args.config_path", args.config_path) predict( @@ -271,5 +272,4 @@ def predict( config_path=args.config_path, prop_val=args.prop_val, background_subs=args.background_subs, - # config_name=args.config_name, )