crfm-helm 0.4.0__py3-none-any.whl → 0.5.10__py3-none-any.whl

helm/clients/audio_language/llama_omni/conversation.py

@@ -0,0 +1,213 @@
+# Adopted from https://github.com/haotian-liu/LLaVA. Below is the original copyright:
+#    Copyright 2023 Haotian Liu
+#
+#    Licensed under the Apache License, Version 2.0 (the "License");
+#    you may not use this file except in compliance with the License.
+#    You may obtain a copy of the License at
+#
+#        http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS,
+#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#    See the License for the specific language governing permissions and
+#    limitations under the License.
+
+import dataclasses
+from enum import auto, Enum
+from typing import List, Any, Union, Optional
+
+
+class SeparatorStyle(Enum):
+    """Different separator style."""
+
+    TWO = auto()
+    PLAIN = auto()
+    LLAMA_2 = auto()
+    LLAMA_3 = auto()
+
+
+@dataclasses.dataclass
+class Conversation:
+    """A class that keeps all conversation history."""
+
+    system: str
+    roles: List[str]
+    messages: List[List[str]]
+    offset: int
+    sep_style: SeparatorStyle = SeparatorStyle.PLAIN
+    sep: str = "###"
+    sep2: str = ""
+    version: str = "Unknown"
+
+    tokenizer_id: str = ""
+    tokenizer: Any = None
+    # Stop criteria (the default one is EOS token)
+    stop_str: Optional[Union[str, List[str]]] = None
+    # Stops generation if meeting any token in this list
+    stop_token_ids: Optional[List[int]] = None
+
+    skip_next: bool = False
+
+    def get_prompt(self):
+        messages = self.messages
+
+        if self.sep_style == SeparatorStyle.TWO:
+            seps = [self.sep, self.sep2]
+            ret = self.system + seps[0]
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message = message[0]
+                    ret += role + ": " + message + seps[i % 2]
+                else:
+                    ret += role + ":"
+        elif self.sep_style == SeparatorStyle.LLAMA_3:
+            wrap_sys = lambda msg: (
+                f"<|start_header_id|>system<|end_header_id|>\n\n{msg}<|eot_id|>" if len(msg) > 0 else msg
+            )
+            ret = "<|begin_of_text|>" + wrap_sys(self.system)
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message = message[0]
+                    ret += f"<|start_header_id|>{role}<|end_header_id|>\n\n"
+                    ret += message.strip() + self.sep2
+                else:
+                    ret += f"<|start_header_id|>{role}<|end_header_id|>\n\n"
+            return ret
+        elif self.sep_style == SeparatorStyle.LLAMA_2:
+            wrap_sys = lambda msg: f"<<SYS>>\n{msg}\n<</SYS>>\n\n" if len(msg) > 0 else msg
+            wrap_inst = lambda msg: f"[INST] {msg} [/INST]"
+            ret = ""
+
+            for i, (role, message) in enumerate(messages):
+                if i == 0:
+                    assert message, "first message should not be none"
+                    assert role == self.roles[0], "first message should come from user"
+                if message:
+                    if type(message) is tuple:
+                        message = message[0]
+                    if i == 0:
+                        message = wrap_sys(self.system) + message
+                    if i % 2 == 0:
+                        message = wrap_inst(message)
+                        ret += self.sep + message
+                    else:
+                        ret += " " + message + " " + self.sep2
+                else:
+                    ret += ""
+            ret = ret.lstrip(self.sep)
+        elif self.sep_style == SeparatorStyle.PLAIN:
+            seps = [self.sep, self.sep2]
+            ret = self.system
+            for i, (role, message) in enumerate(messages):
+                if message:
+                    if type(message) is tuple:
+                        message = message[0]
+                    ret += message + seps[i % 2]
+                else:
+                    ret += ""
+        else:
+            raise ValueError(f"Invalid style: {self.sep_style}")
+
+        return ret
+
+    def append_message(self, role, message):
+        self.messages.append([role, message])
+
+    def to_gradio_chatbot(self):
+        ret = []
+        for i, (role, msg) in enumerate(self.messages[self.offset :]):
+            if i % 2 == 0:
+                if type(msg) is tuple:
+                    msg = msg[0]
+                    ret.append([msg, None])
+                else:
+                    ret.append([msg, None])
+            else:
+                ret[-1][-1] = msg
+        return ret
+
+    def copy(self):
+        return Conversation(
+            system=self.system,
+            roles=self.roles,
+            messages=[[x, y] for x, y in self.messages],
+            offset=self.offset,
+            sep_style=self.sep_style,
+            sep=self.sep,
+            sep2=self.sep2,
+            version=self.version,
+        )
+
+    def dict(self):
+        return {
+            "system": self.system,
+            "roles": self.roles,
+            "messages": self.messages,
+            "offset": self.offset,
+            "sep": self.sep,
+            "sep2": self.sep2,
+        }
+
+
+conv_vicuna_v1 = Conversation(
+    system="A chat between a curious user and an artificial intelligence assistant. "
+    "The assistant gives helpful, detailed, and polite answers to the user's questions.",
+    roles=["USER", "ASSISTANT"],
+    version="v1",
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.TWO,
+    sep=" ",
+    sep2="</s>",
+)
+
+conv_llama_2 = Conversation(
+    system="You are a helpful language and speech assistant. "
+    "You are able to understand the speech content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=["USER", "ASSISTANT"],
+    version="llama_v2",
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_2,
+    sep="<s>",
+    sep2="</s>",
+)
+
+conv_llama_3 = Conversation(
+    system="You are a helpful language and speech assistant. "
+    "You are able to understand the speech content that the user provides, "
+    "and assist the user with a variety of tasks using natural language.",
+    roles=["user", "assistant"],
+    version="llama_v3",
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.LLAMA_3,
+    sep="",
+    sep2="<|eot_id|>",
+)
+
+conv_plain = Conversation(
+    system="",
+    roles=["", ""],
+    messages=[],
+    offset=0,
+    sep_style=SeparatorStyle.PLAIN,
+    sep="</s>",
+)
+
+
+default_conversation = conv_llama_3
+conv_templates = {
+    "v1": conv_vicuna_v1,
+    "plain": conv_plain,
+    "llama_2": conv_llama_2,
+    "llama_3": conv_llama_3,
+}
+
+
+if __name__ == "__main__":
+    print(default_conversation.get_prompt())

helm/clients/audio_language/llama_omni/model/builder.py

@@ -0,0 +1,88 @@
+import os
+
+from transformers import AutoTokenizer, AutoConfig, BitsAndBytesConfig
+import torch
+from helm.clients.audio_language.llama_omni.model.language_model.omni_speech_llama import OmniSpeechLlamaForCausalLM
+from helm.clients.audio_language.llama_omni.model.language_model.omni_speech2s_llama import OmniSpeech2SLlamaForCausalLM
+from helm.clients.audio_language.llama_omni.model.speech_encoder.builder import build_speech_encoder
+
+
+def load_pretrained_model(
+    model_path,
+    model_base,
+    is_lora=False,
+    s2s=False,
+    load_8bit=False,
+    load_4bit=False,
+    device="cuda",
+    use_flash_attn=False,
+    **kwargs,
+):
+    if load_8bit:
+        kwargs["load_in_8bit"] = True
+    elif load_4bit:
+        kwargs["load_in_4bit"] = True
+        kwargs["quantization_config"] = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4",
+        )
+    else:
+        kwargs["torch_dtype"] = torch.float16
+
+    if use_flash_attn:
+        kwargs["attn_implementation"] = "flash_attention_2"
+
+    model_cls = OmniSpeech2SLlamaForCausalLM if s2s else OmniSpeechLlamaForCausalLM
+
+    # Load OmniSpeech model
+    if is_lora:
+        assert model_base is not None, "model_base is required for LoRA models."
+        from language_model.omni_speech_llama import OmniSpeechConfig
+
+        lora_cfg_pretrained = OmniSpeechConfig.from_pretrained(model_path)
+        tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+        print("Loading OmniSpeech from base model...")
+        model = model_cls.from_pretrained(model_base, low_cpu_mem_usage=False, config=lora_cfg_pretrained, **kwargs)
+        print("Loading additional OmniSpeech weights...")
+        if os.path.exists(os.path.join(model_path, "non_lora_trainables.bin")):
+            non_lora_trainables = torch.load(os.path.join(model_path, "non_lora_trainables.bin"), map_location="cpu")
+        non_lora_trainables = {
+            (k[11:] if k.startswith("base_model.") else k): v for k, v in non_lora_trainables.items()
+        }
+        if any(k.startswith("model.model.") for k in non_lora_trainables):
+            non_lora_trainables = {(k[6:] if k.startswith("model.") else k): v for k, v in non_lora_trainables.items()}
+        model.load_state_dict(non_lora_trainables, strict=False)
+
+        from peft import PeftModel
+
+        print("Loading LoRA weights...")
+        model = PeftModel.from_pretrained(model, model_path)
+        print("Merging LoRA weights...")
+        model = model.merge_and_unload()
+        print("Model is loaded...")
+    elif model_base is not None:
+        print("Loading OmniSpeech from base model...")
+        tokenizer = AutoTokenizer.from_pretrained(model_base, use_fast=False)
+        cfg_pretrained = AutoConfig.from_pretrained(model_path)
+        model = model_cls.from_pretrained(model_base, low_cpu_mem_usage=False, config=cfg_pretrained, **kwargs)
+
+        speech_projector_weights = torch.load(os.path.join(model_path, "speech_projector.bin"), map_location="cpu")
+        speech_projector_weights = {k: v.to(torch.float16) for k, v in speech_projector_weights.items()}
+        model.load_state_dict(speech_projector_weights, strict=False)
+        model = model.to(device=device)
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
+        model = model_cls.from_pretrained(model_path, low_cpu_mem_usage=False, **kwargs)
+        model = model.to(device=device)
+
+    model.get_model().speech_encoder = build_speech_encoder(model.config)
+    model.get_model().speech_encoder.to(device=device, dtype=torch.float16)
+
+    if hasattr(model.config, "max_sequence_length"):
+        context_len = model.config.max_sequence_length
+    else:
+        context_len = 2048
+
+    return tokenizer, model, context_len

helm/clients/audio_language/llama_omni/model/language_model/omni_speech2s_llama.py

@@ -0,0 +1,190 @@
+from typing import List, Optional, Tuple, Union, Callable
+
+import torch
+
+from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig
+
+from transformers import PreTrainedModel
+from transformers.generation.streamers import BaseStreamer
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.generation.utils import (
+    GenerationConfig,
+    LogitsProcessorList,
+    StoppingCriteriaList,
+)
+
+from helm.clients.audio_language.llama_omni.model.language_model.omni_speech_llama import OmniSpeechLlamaForCausalLM
+from helm.clients.audio_language.llama_omni.model.speech_generator.builder import build_speech_generator
+from helm.clients.audio_language.llama_omni.model.speech_generator.generation import GenerationWithCTC
+
+
+class OmniSpeech2SConfig(LlamaConfig):
+    model_type = "omni_speech2s_llama"
+
+
+class OmniSpeech2SLlamaForCausalLM(OmniSpeechLlamaForCausalLM, GenerationWithCTC):
+    config_class = OmniSpeech2SConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+        if hasattr(config, "speech_generator_type"):
+            self.speech_generator = build_speech_generator(config)
+
+    def initialize_speech_generator(self, model_args):
+        self.config.speech_generator_type = getattr(model_args, "speech_generator_type", "ctc")
+        self.config.ctc_decoder_config = getattr(model_args, "ctc_decoder_config", "(4,4096,32,11008)")
+        self.config.ctc_upsample_factor = getattr(model_args, "ctc_upsample_factor", 1)
+        self.config.ctc_loss_weight = getattr(model_args, "ctc_loss_weight", 1.0)
+        self.config.unit_vocab_size = getattr(model_args, "unit_vocab_size", 1000)
+        self.tune_speech_generator_only = getattr(model_args, "tune_speech_generator_only", False)
+        if getattr(self, "speech_generator", None) is None:
+            self.speech_generator = build_speech_generator(self.config)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        speech: Optional[torch.FloatTensor] = None,
+        speech_lengths: Optional[torch.LongTensor] = None,
+        tgt_units: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        if inputs_embeds is None:
+            (input_ids, position_ids, attention_mask, past_key_values, inputs_embeds, labels) = (
+                self.prepare_inputs_labels_for_speech_and_text(
+                    input_ids, position_ids, attention_mask, past_key_values, labels, speech, speech_lengths
+                )
+            )
+
+        if self.training:
+            if self.tune_speech_generator_only:
+                with torch.no_grad():
+                    llama_output = super(OmniSpeechLlamaForCausalLM, self).forward(
+                        input_ids=input_ids,
+                        attention_mask=attention_mask,
+                        position_ids=position_ids,
+                        past_key_values=past_key_values,
+                        inputs_embeds=inputs_embeds,
+                        labels=labels,
+                        use_cache=use_cache,
+                        output_attentions=output_attentions,
+                        output_hidden_states=True,
+                        return_dict=return_dict,
+                    )
+                loss = self.speech_generator(llama_output["hidden_states"][-1], labels, tgt_units)
+            else:
+                llama_output = super(OmniSpeechLlamaForCausalLM, self).forward(
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values,
+                    inputs_embeds=inputs_embeds,
+                    labels=labels,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                    output_hidden_states=True,
+                    return_dict=return_dict,
+                )
+                lm_loss = llama_output.loss
+                ctc_loss = self.speech_generator(llama_output["hidden_states"][-1], labels, tgt_units)
+                loss = lm_loss + ctc_loss * self.config.ctc_loss_weight
+        else:
+            llama_output = super(OmniSpeechLlamaForCausalLM, self).forward(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=True,
+                return_dict=return_dict,
+            )
+            loss = llama_output.loss
+
+        return CausalLMOutputWithPast(
+            loss=loss,
+            logits=llama_output.logits,
+            past_key_values=llama_output.past_key_values,
+            hidden_states=llama_output.hidden_states,
+            attentions=llama_output.attentions,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        speech: Optional[torch.Tensor] = None,
+        speech_lengths: Optional[torch.Tensor] = None,
+        generation_config: Optional[GenerationConfig] = None,
+        logits_processor: Optional[LogitsProcessorList] = None,
+        stopping_criteria: Optional[StoppingCriteriaList] = None,
+        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
+        synced_gpus: Optional[bool] = None,
+        assistant_model: Optional["PreTrainedModel"] = None,
+        streamer: Optional["BaseStreamer"] = None,
+        streamer_unit: Optional["BaseStreamer"] = None,
+        streaming_unit_gen=False,
+        negative_prompt_ids: Optional[torch.Tensor] = None,
+        negative_prompt_attention_mask: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        position_ids = kwargs.pop("position_ids", None)
+        attention_mask = kwargs.pop("attention_mask", None)
+        if "inputs_embeds" in kwargs:
+            raise NotImplementedError("`inputs_embeds` is not supported")
+
+        if speech is not None:
+            (inputs, position_ids, attention_mask, _, inputs_embeds, _) = (
+                self.prepare_inputs_labels_for_speech_and_text(
+                    inputs, position_ids, attention_mask, None, None, speech, speech_lengths
+                )
+            )
+        else:
+            inputs_embeds = self.get_model().embed_tokens(inputs)
+        outputs = GenerationWithCTC.generate(
+            self,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=True,
+            return_dict_in_generate=True,
+            streaming_unit_gen=streaming_unit_gen,
+            **kwargs,
+        )
+
+        hidden_states = outputs["hidden_states"]
+        hidden_states = torch.cat(
+            [hidden_states[0][-1][:, -1:, :]] + [hidden_states[i][-1] for i in range(1, len(hidden_states))], dim=1
+        )
+        ctc_pred = self.speech_generator.predict(hidden_states.squeeze(0))
+
+        return outputs.sequences, ctc_pred
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        speech = kwargs.pop("speech", None)
+        speech_lengths = kwargs.pop("speech_lengths", None)
+        inputs = super().prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
+        )
+        if speech is not None:
+            inputs["speech"] = speech
+            inputs["speech_lengths"] = speech_lengths
+        return inputs
+
+
+AutoConfig.register("omni_speech2s_llama", OmniSpeech2SConfig)
+AutoModelForCausalLM.register(OmniSpeech2SConfig, OmniSpeech2SLlamaForCausalLM)

helm/clients/audio_language/llama_omni/model/language_model/omni_speech_llama.py

@@ -0,0 +1,118 @@
+from typing import List, Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+
+from transformers import AutoConfig, AutoModelForCausalLM, LlamaConfig, LlamaModel, LlamaForCausalLM
+
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.generation.utils import GenerateOutput
+
+from helm.clients.audio_language.llama_omni.model.omni_speech_arch import OmniSpeechMetaModel, OmniSpeechMetaForCausalLM
+
+
+class OmniSpeechConfig(LlamaConfig):
+    model_type = "omni_speech_llama"
+
+
+class OmniSpeechLlamaModel(OmniSpeechMetaModel, LlamaModel):
+    config_class = OmniSpeechConfig
+
+    def __init__(self, config: LlamaConfig):
+        super(OmniSpeechLlamaModel, self).__init__(config)
+
+
+class OmniSpeechLlamaForCausalLM(LlamaForCausalLM, OmniSpeechMetaForCausalLM):
+    config_class = OmniSpeechConfig
+
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = OmniSpeechLlamaModel(config)
+        self.pretraining_tp = config.pretraining_tp
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_model(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[List[torch.FloatTensor]] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        speech: Optional[torch.FloatTensor] = None,
+        speech_lengths: Optional[torch.LongTensor] = None,
+        tgt_units: Optional[torch.LongTensor] = None,
+        return_dict: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+    ) -> Union[Tuple, CausalLMOutputWithPast]:
+
+        if inputs_embeds is None:
+            (input_ids, position_ids, attention_mask, past_key_values, inputs_embeds, labels) = (
+                self.prepare_inputs_labels_for_speech_and_text(
+                    input_ids, position_ids, attention_mask, past_key_values, labels, speech, speech_lengths
+                )
+            )
+
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+    @torch.no_grad()
+    def generate(
+        self,
+        inputs: Optional[torch.Tensor] = None,
+        speech: Optional[torch.Tensor] = None,
+        speech_lengths: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> Union[GenerateOutput, torch.LongTensor]:
+        position_ids = kwargs.pop("position_ids", None)
+        attention_mask = kwargs.pop("attention_mask", None)
+        if "inputs_embeds" in kwargs:
+            raise NotImplementedError("`inputs_embeds` is not supported")
+
+        if speech is not None:
+            (inputs, position_ids, attention_mask, _, inputs_embeds, _) = (
+                self.prepare_inputs_labels_for_speech_and_text(
+                    inputs, position_ids, attention_mask, None, None, speech, speech_lengths
+                )
+            )
+        else:
+            inputs_embeds = self.get_model().embed_tokens(inputs)
+
+        return super().generate(
+            position_ids=position_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, **kwargs
+        )
+
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):
+        speech = kwargs.pop("speech", None)
+        speech_lengths = kwargs.pop("speech_lengths", None)
+        inputs = super().prepare_inputs_for_generation(
+            input_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, **kwargs
+        )
+        if speech is not None:
+            inputs["speech"] = speech
+            inputs["speech_lengths"] = speech_lengths
+        return inputs
+
+
+AutoConfig.register("omni_speech_llama", OmniSpeechConfig)
+AutoModelForCausalLM.register(OmniSpeechConfig, OmniSpeechLlamaForCausalLM)