Bugfix/invalid structures

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"]
     )

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,
     )