Feature/loader factory

atomgpt/inverse_models/dataset_utils.py

@@ -6,6 +6,16 @@
 
 from typing import Union, Callable, Optional, List, Dict
 import torch
+from typing import Any
+from jarvis.core.atoms import Atoms
+from jarvis.io.vasp.inputs import Poscar
+from jarvis.core.composition import Composition
+from atomgpt.inverse_models.utils import (
+    gen_atoms,
+    text2atoms,
+    get_crystal_string_t,
+    get_figlet,
+)
 
 
 # From https://www.geeksforgeeks.org/longest-common-substring-array-strings/
@@ -753,6 +763,100 @@ def _tokenize(example):
         )
     pass
     return dataset
+pass
 
+def get_input(config=None, chem="", val=10):
+    if config.chem_info == "none":
+        prefix = ""
+    elif config.chem_info == "element_list":
+        prefix = (
+            "The chemical elements are "
+            + chem  # atoms.composition.search_string
+            + " . "
+        )
+    elif config.chem_info == "element_dict":
+        prefix = (
+            "The chemical contents are "
+            + chem  # atoms.composition.search_string
+            + " . "
+        )
+    elif config.chem_info == "formula":
+        prefix = (
+            "The chemical formula is "
+            + chem  # atoms.composition.reduced_formula
+            + " . "
+        )
 
-pass
+    inp = (
+        prefix
+        + "The  "
+        + config.prop
+        + " is "
+        + str(val)
+        + "."
+        + config.output_prompt
+    )
+    return inp
+
+def make_alpaca_json(
+    dataset=[],
+    jids=[],
+    # prop="Tc_supercon",
+    # instruction="",
+    include_jid=False,
+    # chem_info="",
+    # output_prompt="",
+    config=None,
+):
+    mem = []
+    print("config.prop", config.prop)
+    for i in dataset:
+        if i[config.prop] != "na" and i[config.id_tag] in jids:
+            atoms = Atoms.from_dict(i["atoms"])
+            info = {}
+            if include_jid:
+                info["id"] = i[config.id_tag]
+            info["instruction"] = config.instruction
+            if config.chem_info == "none":
+                chem = ""
+            elif config.chem_info == "element_list":
+                chem = atoms.composition.search_string
+            elif config.chem_info == "element_dict":
+                comp = Composition.from_string(
+                    atoms.composition.reduced_formula
+                )
+                chem = comp.to_dict()
+                chem = str(dict(sorted(chem.items())))
+            elif config.chem_info == "formula":
+                chem = atoms.composition.reduced_formula
+
+            inp = get_input(config=config, val=i[config.prop], chem=chem)
+            info["input"] = inp
+
+            info["output"] = get_crystal_string_t(atoms)
+            mem.append(info)
+    return mem
+
+def alpaca_formatting_prompts_func(examples: Dict[str, Any], alpaca_prompt: str, eos_token: str) -> Dict[str, List[str]]:
+    inst = examples["instruction"]
+    inp  = examples["input"]
+    out  = examples["output"]
+    texts = [alpaca_prompt.format(i, x, y) + eos_token for i, x, y in zip(inst, inp, out)]
+    return {"text": texts}
+
+def harmony_formatting_prompts_func(examples: Dict[str, Any], tokenizer) -> Dict[str, List[str]]:
+    inst = examples["instruction"]
+    inp  = examples["input"]
+    out  = examples["output"]
+    texts: List[str] = []
+    for i, x, y in zip(inst, inp, out):
+        messages = []
+        i = (i or "").strip()
+        x = (x or "").strip()
+        y = (y or "").strip()
+        if i:
+            messages.append({"role": "developer", "content": i})
+        messages.append({"role": "user", "content": x})
+        messages.append({"role": "assistant", "content": y})
+        texts.append(tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=False))
+    return {"text": texts}

atomgpt/inverse_models/factories.py

@@ -0,0 +1,149 @@
+# factories.py
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from atomgpt.inverse_models.products import LoadedModel
+from typing import Callable
+from typing import TYPE_CHECKING
+if TYPE_CHECKING:
+    from atomgpt.inverse_models.inverse_models import TrainingPropConfig
+from peft import PeftModel
+from typing import Dict
+from atomgpt.inverse_models.dataset_utils import alpaca_formatting_prompts_func
+from atomgpt.inverse_models.dataset_utils import harmony_formatting_prompts_func
+from functools import partial
+from typing import List
+
+
+class LanguageModelFactory(ABC):
+    @abstractmethod
+    def load_for_training(self, config: TrainingPropConfig) -> LoadedModel:
+        pass
+
+    @abstractmethod
+    def load_for_inference(self, checkpoint_path: str, config: TrainingPropConfig) -> LoadedModel:
+        pass
+
+    @abstractmethod
+    def get_formatting_prompts_func(self, config, model, tokenizer) -> Callable:
+        pass
+
+
+class AtomGPTFactory(LanguageModelFactory):
+    def load_for_training(self, config: TrainingPropConfig) -> LoadedModel:
+        from atomgpt.inverse_models.loader import FastLanguageModel as AtomGPTFastLanguageModel
+        model, tokenizer = AtomGPTFastLanguageModel.from_pretrained(
+            model_name=config.model_name,
+            max_seq_length=config.max_seq_length,
+            dtype=config.dtype,
+            load_in_4bit=config.load_in_4bit
+        )
+        if not isinstance(model, PeftModel):
+            # import sys
+            print("Not yet a peft model, converting into peft model")
+            # sys.exit()
+            model = AtomGPTFastLanguageModel.get_peft_model(
+                model,
+                r=config.lora_rank,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+                target_modules=[
+                    "q_proj",
+                    "k_proj",
+                    "v_proj",
+                    "o_proj",
+                    "gate_proj",
+                    "up_proj",
+                    "down_proj",
+                ],
+                lora_alpha=config.lora_alpha,
+                lora_dropout=0,  # Supports any, but = 0 is optimized
+                bias="none",  # Supports any, but = "none" is optimized
+                use_gradient_checkpointing=True,
+                random_state=3407,
+                use_rslora=False,  # We support rank stabilized LoRA
+                loftq_config=None,  # And LoftQ
+            )
+            print("Peft model created")
+        EOS_TOKEN = tokenizer.eos_token
+        return LoadedModel(model=model, tokenizer=tokenizer)
+
+    def load_for_inference(self, checkpoint_path: str, config: TrainingPropConfig) -> LoadedModel:
+        model, tokenizer = FastLanguageModel.from_pretrained(
+            model_name=checkpoint_path,
+            max_seq_length=config.max_seq_length,
+            dtype=config.dtype,
+            load_in_4bit=config.load_in_4bit,
+        )
+        AtomGPTFastLanguageModel.for_inference(model)
+        return LoadedModel(model=model, tokenizer=tokenizer)
+
+    def get_formatting_prompts_func(self, config, model, tokenizer) -> Callable:
+        eos = tokenizer.eos_token or "</s>"
+        return partial(alpaca_formatting_prompts_func, alpaca_prompt=config.alpaca_prompt, eos_token=eos)
+
+
+class GPTOSSFactory(LanguageModelFactory):
+    def load_for_training(self, config: TrainingPropConfig) -> LoadedModel:
+        from unsloth import FastLanguageModel as UnslothFastLanguageModel
+        model, tokenizer = UnslothFastLanguageModel.from_pretrained(
+            model_name=config.model_name,
+            max_seq_length=config.max_seq_length,
+            dtype=config.dtype,
+            load_in_4bit=config.load_in_4bit,
+            full_finetuning = False,
+        )
+        if not isinstance(model, PeftModel):
+            print("Not yet a peft model, converting into peft model")
+            model = UnslothFastLanguageModel.get_peft_model(
+                model,
+                r=config.lora_rank,  # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+                target_modules=[
+                    "q_proj",
+                    "k_proj",
+                    "v_proj",
+                    "o_proj",
+                    "gate_proj",
+                    "up_proj",
+                    "down_proj",
+                ],
+                lora_alpha=config.lora_alpha,
+                lora_dropout=0,  # Supports any, but = 0 is optimized
+                bias="none",  # Supports any, but = "none" is optimized
+                use_gradient_checkpointing=True,
+                random_state=3407,
+                use_rslora=False,  # We support rank stabilized LoRA
+                loftq_config=None,  # And LoftQ
+            )
+            print("Peft model created")
+        return LoadedModel(model=model, tokenizer=tokenizer)
+
+    def load_for_inference(self, checkpoint_path: str, config: TrainingPropConfig) -> LoadedModel:
+        model, tokenizer = UnslothFastLanguageModel.from_pretrained(
+            model_name=checkpoint_path,
+            max_seq_length=config.max_seq_length,
+            dtype=config.dtype,
+            load_in_4bit=config.load_in_4bit,
+        )
+        UnslothFastLanguageModel.for_inference(model)
+        return LoadedModel(model=model, tokenizer=tokenizer)
+
+    def get_formatting_prompts_func(self, config, model, tokenizer) -> Callable:
+        return partial(harmony_formatting_prompts_func, tokenizer=tokenizer)
+
+FACTORY_REGISTRY: Dict[str, type[LanguageModelFactory]] = {
+    "gemma": AtomGPTFactory,
+    "qwen": AtomGPTFactory,
+    "Meta": AtomGPTFactory,
+    "Llama": AtomGPTFactory,
+    "llama": AtomGPTFactory,
+    "Mistral": AtomGPTFactory,
+    "mistral": AtomGPTFactory,
+    "gpt-oss": GPTOSSFactory,
+}
+
+def get_lm_factory(config: TrainingPropConfig) -> LanguageModelFactory:
+    model_name = config.model_name
+    if "gpt-oss" in model_name:
+        return GPTOSSFactory()
+    else:
+        return AtomGPTFactory()