feat(components): Update policy to reward model name mapping in funct…

…ion based component in _implementation/llm PiperOrigin-RevId: 560479276
kubeflow · Aug 27, 2023 · 110e082 · 110e082
1 parent 63a0803
commit 110e082
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/...nents/google-cloud/google_cloud_pipeline_components/_implementation/llm/function_based.py b/...nents/google-cloud/google_cloud_pipeline_components/_implementation/llm/function_based.py
@@ -294,15 +294,15 @@ def resolve_reference_model_metadata(
       'llama-2-7b-chat': reference_model_metadata(
           large_model_reference='LLAMA_2_7B_CHAT',
           reference_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_7b_chat/',
-          reward_model_reference='LLAMA_2_7B_CHAT',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_7b_chat/',
+          reward_model_reference='LLAMA_2_7B',
+          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_7b/',
           is_supported=True,
       ),
       'llama-2-13b-chat': reference_model_metadata(
           large_model_reference='LLAMA_2_13B_CHAT',
           reference_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b_chat/',
-          reward_model_reference='LLAMA_2_13B_CHAT',
-          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b_chat/',
+          reward_model_reference='LLAMA_2_13B',
+          reward_model_path='gs://vertex-rlhf-restricted/pretrained_models/llama/t5x_llama_2_13b/',
           is_supported=True,
       ),
   }