Spectrum tuning integration

CHANGELOG.md

@@ -12,6 +12,7 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 * Added option for ICL example choice ordering - fixed, swapped, or random
 * Added option to swap choice ordering for comparative regression pipeline ADM component
 * Added capability to resolve pipeline ADM step output conflicts with a custom function
+* Added inference engine for spectrum tuned LLMs that appropriately reformats chat template roles
 
 ## 0.5.9
 

align_system/algorithms/outlines_inference_engine.py

@@ -150,3 +150,57 @@ def _sampler_repr(sampler):
                        model_kwargs={self.model_kwargs},
                        tokenizer_kwargs={self.tokenizer_kwargs},
                        )""").strip()
+
+
+class SpectrumTunedInferenceEngine(OutlinesTransformersInferenceEngine):
+    def __init__(self,
+                 model_name,
+                 device='auto',
+                 precision='full',
+                 max_generator_tokens=None,
+                 sampler=MultinomialSampler(),
+                 inference_batch_size=5,
+                 model_kwargs={},
+                 tokenizer_kwargs={}):
+        super().__init__(model_name,
+                         device,
+                         precision,
+                         max_generator_tokens,
+                         sampler,
+                         inference_batch_size,
+                         model_kwargs,
+                         tokenizer_kwargs)
+
+    def dialog_to_prompt(self, dialog):
+        tokenizer = self.model.tokenizer.tokenizer
+
+        # Use roles spectrum tuned models expect
+        # https://github.com/tsor13/spectrum/blob/main/README.md
+        for element in dialog:
+            if element.role == "system":
+                element.role = "description"
+            elif element.role == "user":
+                element.role = "input"
+            elif element.role == "assistant":
+                element.role = "output"
+            else:
+                raise RuntimeError(f"{element.role} dialog element unrecognized.")
+
+        try:
+            encoded_dialog = tokenizer.apply_chat_template(dialog)
+        except jinja2.exceptions.TemplateError:
+            # Assume that the tokenizer chat template doesn't accept
+            # system messages; combine system message first user
+            # message
+            # Ensure each dialog element is a dict
+            system_msg, user_msg, *rest = [dict(d) for d in dialog]
+
+            assert user_msg['role'] == 'user'
+
+            updated_content = system_msg['content'] + '\n' + user_msg['content']
+
+            dialog = [{'role': 'user', 'content': updated_content}, *rest]
+
+            encoded_dialog = tokenizer.apply_chat_template(dialog)
+
+        return tokenizer.decode(encoded_dialog)

align_system/configs/adm/phase2_pipeline_fewshot_comparative_regression_spectrum_tuned.yaml

@@ -0,0 +1,59 @@
+name: phase2_pipeline_zeroshot_comparative_regression_spectrum_tuned
+
+defaults:
+  # Import defaults into this namspace (adm) as @name, for further
+  # customization
+
+  # Shared variables / components
+  - /attribute@mu: medical_urgency
+  - /attribute@af: affiliation_focus
+  - /attribute@mf: merit_focus
+  - /attribute@ss: search_or_stay
+  - /attribute@ps: personal_safety
+  - /inference_engine@structured_inference_engine: spectrum_tuned_greedy
+  - /template/scenario_description@scenario_description_template: phase2
+  - /template/prompt@prompt_template: phase2_comparative_regression
+  - /template/output_schema@comparative_regression_choice_schema: phase2_comparative_regression_choice
+  # ADM components to be used in "steps"
+  - /adm_component/misc@step_definitions.format_choices: itm_format_choices
+  - /adm_component/icl@step_definitions.regression_icl: phase2_comparative
+  - /adm_component/regression@step_definitions.comparative_regression: phase2_comparative_no_template
+  - /adm_component/misc@step_definitions.regression_rule_based_correction: phase2_regression_rule_based_correction
+  - /adm_component/alignment@step_definitions.scalar_alignment: medical_urgency_scalar
+  - /adm_component/misc@step_definitions.justification_from_reasonings: justification_from_reasonings
+  - /adm_component/misc@step_definitions.ensure_chosen_action: ensure_chosen_action
+  - /adm_component/misc@step_definitions.populate_choice_info: populate_choice_info
+  # Use definitions in this file to override defaults defined above
+  - _self_
+
+attribute_definitions:
+  medical: ${adm.mu}
+  affiliation: ${adm.af}
+  merit: ${adm.mf}
+  search: ${adm.ss}
+  personal_safety: ${adm.ps}
+
+step_definitions:
+  regression_icl:
+    scenario_description_template: ${ref:adm.scenario_description_template}
+    attributes: ${adm.attribute_definitions}
+    prompt_template: ${ref:adm.prompt_template}
+
+  comparative_regression:
+    scenario_description_template: ${ref:adm.scenario_description_template}
+    prompt_template: ${ref:adm.prompt_template}
+    score_schema_template: ${adm.comparative_regression_choice_schema}
+
+instance:
+  _target_: align_system.algorithms.pipeline_adm.PipelineADM
+
+  steps:
+    # Reference the step instances we want to use in order
+    - ${ref:adm.step_definitions.format_choices}
+    - ${ref:adm.step_definitions.regression_icl}
+    - ${ref:adm.step_definitions.comparative_regression}
+    - ${ref:adm.step_definitions.regression_rule_based_correction}
+    - ${ref:adm.step_definitions.scalar_alignment}
+    - ${ref:adm.step_definitions.justification_from_reasonings}
+    - ${ref:adm.step_definitions.ensure_chosen_action}
+    - ${ref:adm.step_definitions.populate_choice_info}

align_system/configs/inference_engine/spectrum_tuned_greedy.yaml

@@ -0,0 +1,6 @@
+_target_: align_system.algorithms.outlines_inference_engine.SpectrumTunedInferenceEngine
+
+model_name: tsor13/spectrum-Llama-3.1-8B-v1
+precision: half
+sampler:
+  _target_: outlines.samplers.GreedySampler