Merge pull request caikit#314 from gkumbhat/add_granite_modeling_llam…

…a_main

Add granite modeling llama main

caikit_nlp/modules/text_generation/peft_prompt_tuning.py

@@ -440,7 +440,7 @@ def train(
 
         # Remove _name_or_path field as a model can be
         # saved in different location but still same
-        del base_model_config["_name_or_path"]
+        base_model_config.pop("_name_or_path", None)
         error.value_check(
             "<NLP07232147E>",
             "_name_or_path" not in base_model_config,
@@ -585,7 +585,8 @@ def load(
             if peft_config.task_type == "CAUSAL_LM":
                 # get the transformers Causal LM model
                 base_model = AutoModelForCausalLM.from_pretrained(
-                    peft_config.base_model_name_or_path
+                    peft_config.base_model_name_or_path,
+                    torch_dtype=torch_dtype,
                 )
                 # get the PEFT causal LM model
                 model = PeftModel.from_pretrained(base_model, model_config)

caikit_nlp/resources/pretrained_model/hf_auto_causal_lm.py

@@ -35,6 +35,11 @@
 
 # Local
 from ...data_model import GenerationTrainRecord, PromptOutputModelType
+
+# Note: Below module is imported to allow loading of fm stack sphinx models
+from ...toolkit.text_generation import (  # pylint: disable=unused-import
+    granite_modeling_llama,
+)
 from ...toolkit.verbalizer_utils import render_verbalizer
 from .base import PretrainedModelBase