meetkai/functionary-medium-v3.0

1# coding=utf-82# Copyright (c) 2024, MeetKai Inc. All rights reserved.3"""PyTorch LLaMA model."""4 5import json6from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Union7 8import torch9import torch.utils.checkpoint10 11from transformers.generation.configuration_utils import GenerationConfig12from transformers.generation.logits_process import LogitsProcessorList13from transformers.generation.stopping_criteria import StoppingCriteriaList14from transformers.generation.utils import (15    GenerateBeamDecoderOnlyOutput,16    GenerateBeamEncoderDecoderOutput,17    GenerateDecoderOnlyOutput,18    GenerateEncoderDecoderOutput19)20from transformers.models.llama.modeling_llama import LlamaForCausalLM21from transformers.utils import logging22 23 24if TYPE_CHECKING:25    from transformers.modeling_utils import PreTrainedModel26    from transformers.generation.streamers import BaseStreamer27 28logger = logging.get_logger(__name__)29 30GenerateNonBeamOutput = Union[GenerateDecoderOnlyOutput, GenerateEncoderDecoderOutput]31GenerateBeamOutput = Union[GenerateBeamDecoderOnlyOutput, GenerateBeamEncoderDecoderOutput]32GenerateOutput = Union[GenerateNonBeamOutput, GenerateBeamOutput]33 34 35class FunctionaryForCausalLM(LlamaForCausalLM):36 37    def generate_tool_use(38        self,39        inputs: Optional[torch.Tensor] = None,40        generation_config: Optional[GenerationConfig] = None,41        logits_processor: Optional[LogitsProcessorList] = None,42        stopping_criteria: Optional[StoppingCriteriaList] = None,43        prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,44        synced_gpus: Optional[bool] = None,45        assistant_model: Optional["PreTrainedModel"] = None,46        streamer: Optional["BaseStreamer"] = None,47        negative_prompt_ids: Optional[torch.Tensor] = None,48        negative_prompt_attention_mask: Optional[torch.Tensor] = None,49        **kwargs,50    ) -> Union[GenerateOutput, torch.LongTensor]:51 52        tokenizer = kwargs.pop("tokenizer", None)  # Pull this out first, we use it to parse raw output53 54        results = self.generate(55            inputs=inputs,56            generation_config=generation_config,57            logits_processor=logits_processor,58            stopping_criteria=stopping_criteria,59            prefix_allowed_tokens_fn=prefix_allowed_tokens_fn,60            synced_gpus=synced_gpus,61            assistant_model=assistant_model,62            streamer=streamer,63            negative_prompt_ids=negative_prompt_ids,64            negative_prompt_attention_mask=negative_prompt_attention_mask,65            **kwargs,66        )67 68        input_ids = kwargs.pop("input_ids")69        function_call_token = ">>>"70        71        correct_results = []72        for input_id, result in zip(input_ids, results):73            final_output_json = {"role": "assistant", "content": None, "tool_calls": None}74            tool_calls = []75            raw_output_str = tokenizer.decode(result[len(input_id):].cpu())76            chunks = raw_output_str.split(function_call_token)77            for i, chunk in enumerate(chunks):78                if len(chunk) == 0:79                    continue80                81                name, arguments = chunk[: chunk.index("\n")], chunk[chunk.index("\n") + 1:]82                arguments = arguments[: -len("<|eot_id|>")] if arguments.endswith("<|eot_id|>") else arguments83                84                if name == "all":85                    if final_output_json["content"] is None:86                        final_output_json["content"] = [arguments]87                    else:88                        final_output_json["content"].append(arguments)89                else:90                    tool_calls.append(91                        {92                            "name": name,93                            "arguments": arguments94                        }95                    )96 97            if final_output_json["content"] is not None:98                final_output_json["content"] = " ".join(final_output_json["content"])99            if len(tool_calls) > 0:100                final_output_json["tool_calls"] = tool_calls101            final_output_str = json.dumps(final_output_json, indent=4)102            final_output_ids = tokenizer(final_output_str, add_special_tokens=False)["input_ids"]103            correct_results.append(104                torch.cat(105                    (result[:len(input_id)].cpu(), torch.tensor(final_output_ids))106                )107            )108        max_len = max([tensor.shape[0] for tensor in correct_results])109        correct_results = [110            torch.nn.functional.pad(111                correct_result, (0, max_len - correct_result.shape[0]), value=tokenizer.eos_token_id112            ) for correct_result in correct_results113        ]114        correct_results = torch.stack(correct_results)115        116        return correct_results