Upload folder using huggingface_hub

configuration_kraken.py

@@ -0,0 +1,8 @@
+from transformers import PretrainedConfig
+
+class KrakenConfig(PretrainedConfig):
+    model_type = "kraken"
+    
+    def __init__(self, config_dict=None, **kwargs):
+        super().__init__(**kwargs)
+        self.config_dict = config_dict or {}

kraken_model/config.json

@@ -0,0 +1,40 @@
+{
+  "architectures": [
+    "KrakenForCausalLM"
+  ],
+  "config_dict": {
+    "class_indices": {
+      "LABEL_0": 0,
+      "LABEL_1": 1,
+      "LABEL_2": 2,
+      "LABEL_3": 3,
+      "LABEL_4": 4
+    },
+    "model_type": "kraken",
+    "models": {
+      "expert1": "microsoft/Phi-3-medium-128k-instruct",
+      "expert2": "gorilla-llm/gorilla-openfunctions-v2",
+      "expert3": "ise-uiuc/Magicoder-S-DS-6.7B",
+      "expert4": "defog/llama-3-sqlcoder-8b",
+      "expert5": "VAGOsolutions/Llama-3-SauerkrautLM-8b-Instruct"
+    },
+    "quantization": {
+      "expert1": null,
+      "expert2": null,
+      "expert3": null,
+      "expert4": null,
+      "expert5": null
+    },
+    "router": "./kraken_router",
+    "tokenizers": {
+      "expert1": "microsoft/Phi-3-medium-128k-instruct",
+      "expert2": "gorilla-llm/gorilla-openfunctions-v2",
+      "expert3": "ise-uiuc/Magicoder-S-DS-6.7B",
+      "expert4": "defog/llama-3-sqlcoder-8b",
+      "expert5": "VAGOsolutions/Llama-3-SauerkrautLM-8b-Instruct"
+    }
+  },
+  "model_type": "kraken",
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.0"
+}

kraken_model/generation_config.json

@@ -0,0 +1,4 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.41.0"
+}

kraken_model/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:e0acccf86abbb885412274d405d4555248636fe829dedbd7655bceb29a023535
+size 1856007992

kraken_router/added_tokens.json

@@ -0,0 +1,5 @@
+{
+  "<|endoftext|>": 151643,
+  "<|im_end|>": 151645,
+  "<|im_start|>": 151644
+}

kraken_router/config.json

@@ -0,0 +1,44 @@
+{
+  "_name_or_path": "Qwen/Qwen1.5-0.5B",
+  "architectures": [
+    "Qwen2ForSequenceClassification"
+  ],
+  "attention_dropout": 0.0,
+  "bos_token_id": 151643,
+  "eos_token_id": 151643,
+  "hidden_act": "silu",
+  "hidden_size": 1024,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1",
+    "2": "LABEL_2",
+    "3": "LABEL_3",
+    "4": "LABEL_4"
+  },
+  "initializer_range": 0.02,
+  "intermediate_size": 2816,
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1,
+    "LABEL_2": 2,
+    "LABEL_3": 3,
+    "LABEL_4": 4
+  },
+  "max_position_embeddings": 32768,
+  "max_window_layers": 21,
+  "model_type": "qwen2",
+  "num_attention_heads": 16,
+  "num_hidden_layers": 24,
+  "num_key_value_heads": 16,
+  "pad_token_id": 151643,
+  "problem_type": "single_label_classification",
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 1000000.0,
+  "sliding_window": 32768,
+  "tie_word_embeddings": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.41.0",
+  "use_cache": true,
+  "use_sliding_window": false,
+  "vocab_size": 151936
+}

kraken_router/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0462ee8274d9aab1297d68d7363e40940d90fd274cb2679c4790fdf4371f2808
+size 1856004208

kraken_router/special_tokens_map.json

@@ -0,0 +1,14 @@
+{
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>"
+  ],
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": "<|endoftext|>"
+}

kraken_router/tokenizer_config.json

@@ -0,0 +1,43 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "151643": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151644": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "151645": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>"
+  ],
+  "bos_token": null,
+  "chat_template": "{% for message in messages %}{% if loop.first and messages[0]['role'] != 'system' %}{{ '<|im_start|>system\nYou are a helpful assistant<|im_end|>\n' }}{% endif %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "model_max_length": 32768,
+  "pad_token": "<|endoftext|>",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null
+}

modeling_kraken.py

@@ -0,0 +1,82 @@
+import torch
+from transformers import PreTrainedModel, AutoTokenizer, AutoModelForCausalLM, AutoModelForSequenceClassification, TextClassificationPipeline
+from configuration_kraken import KrakenConfig
+import tokenizer_template_switch
+
+class KrakenForCausalLM(PreTrainedModel):
+    config_class = KrakenConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.tokenizers = {key: AutoTokenizer.from_pretrained(name, device_map="auto") for key, name in config.config_dict['tokenizers'].items()}
+        self.models = self.load_expert_models(config.config_dict['models'], config.config_dict['quantization'])
+        self.router_model = AutoModelForSequenceClassification.from_pretrained(config.config_dict['router'], trust_remote_code=True,device_map="auto")
+        self.tokenizer = AutoTokenizer.from_pretrained(config.config_dict['router'], trust_remote_code=True,device_map="auto")
+        self.router = TextClassificationPipeline(model=self.router_model, tokenizer=self.tokenizer)
+        self.models_indices = config.config_dict['class_indices']
+
+    def load_expert_models(self, models_dict, quantization_dict):
+        models = {}
+        for key, name in models_dict.items():
+            quantization = quantization_dict.get(key)
+            if quantization == "8bit":
+                models[key] = AutoModelForCausalLM.from_pretrained(name, trust_remote_code=True, device_map="auto", load_in_8bit=True, torch_dtype="auto")
+            elif quantization == "4bit":
+                models[key] = AutoModelForCausalLM.from_pretrained(name, trust_remote_code=True, device_map="auto", load_in_4bit=True, torch_dtype="auto")
+            elif quantization == "awq":
+                models[key] = self.load_awq_model(name)
+            else:
+                models[key] = AutoModelForCausalLM.from_pretrained(name, trust_remote_code=True, device_map="auto", torch_dtype="auto")
+        return models
+
+    def load_awq_model(self, name):
+        return AutoModelForCausalLM.from_pretrained(name, trust_remote_code=True, device_map="auto")
+
+    def tokenize_inputs(self, text, model_key):
+        return self.tokenizers[model_key](text, return_tensors="pt")
+
+    def determine_model(self, text):
+        prediction = self.router(text)[0]["label"]
+        model_decision_index = self.models_indices[prediction]
+        model_keys = ['expert1', 'expert2', 'expert3', 'expert4','expert5']
+        return model_keys[model_decision_index]
+    
+    def expert_tokenizer(self, text):
+        model_key = self.determine_model(text)
+        return self.tokenizers[model_key]
+
+
+    def generate(self, input_ids, **generate_kwargs):
+        # Tokenize the input_ids
+        text = self.tokenizer.batch_decode(input_ids, skip_special_tokens=False)[0]
+
+        msgs = tokenizer_template_switch.recover_chat_messages(text, self.tokenizer)
+        if msgs and msgs[0]['role'] == 'system' and msgs[0]['content']=='<|im_start|>system':
+            # Delete the first element
+            msgs.pop(0)  
+        # Check if the last element has the role 'assistant'
+        if msgs and msgs[-1]['role'] == 'assistant':
+            # Delete the last element
+            msgs.pop()  
+
+        # Determine the model key using the existing routing logic
+        model_key = self.determine_model(text)
+        # Show the routing result
+        print(f"Choosing {model_key} ..")
+        # Retrieve the model from the dictionary
+        model = self.models[model_key]
+
+        mod_txt = self.tokenizers[model_key].apply_chat_template(msgs, tokenize=False, add_generation_prompt=True)
+        current_device = input_ids.device if isinstance(input_ids, torch.Tensor) else 'cpu'
+        
+        # Tokenize accordingly to the best model
+
+        tok = self.tokenizers[model_key](mod_txt, return_tensors="pt")
+        tok_input_ids = tok.input_ids.to(current_device)  
+        tok_attention_mask = tok.attention_mask.to(current_device)
+
+        # Generate text using the retrieved model
+        return model.generate(tok_input_ids, attention_mask=tok_attention_mask, **generate_kwargs)
+    
+
+      

tokenizer_template_switch.py

@@ -0,0 +1,95 @@
+import re
+from transformers import AutoTokenizer
+
+def extract_separators(template):
+    """
+    Extracts separators used in the tokenization template.
+    """
+    # Adjust the regex to correctly match the specific pattern between '{{' and '+ message["content"] +'
+    pattern = r"\{\{\s*([^{}]+?)\s*\+ message\['content'\]"
+    matches = re.findall(pattern, template)
+    # Clean up any extra spaces and return the matches
+    separators = [match.strip() for match in matches]
+
+    if any("message['role']" in element for element in separators):
+        roles = ["system", "user", "assistant"]
+        separators_ = []
+        for role in roles:
+            separators_.append(separators[0].replace(" + message['role'] + ", role).replace("'",""))
+        return separators_
+
+    return separators
+
+def detect_eos_token(jinja_template, tokenizer):
+    if "<|im_end|>" in jinja_template:
+        return "<|im_end|>"
+    if "</s>" in jinja_template:
+        return "</s>"
+    if "eos_token" in jinja_template:
+        return tokenizer.eos_token    
+    else:
+        return "<|endoftext|>"
+
+def recover_messages(formatted_message, separators, eos_token):
+    """
+    Recovers the original messages from the formatted message string.
+    """
+    # Split the formatted message using the end-of-string token
+    split_messages = formatted_message.split(eos_token)
+    
+    # Remove the last empty string if it exists due to a trailing separator
+    if split_messages and split_messages[-1].strip() == '':
+        split_messages.pop()
+
+    # Prepare the list to hold the recovered messages
+    recovered_messages = []
+
+    # Define roles after the first message, alternating between "user" and "assistant"
+    alternate_roles = ["user", "assistant"]
+    
+    # Iterate over the split messages
+    for index, message_content in enumerate(split_messages):
+        # Determine the role, starting with "system" for the first message
+        # then alternating between "user" and "assistant" for subsequent messages
+        if index == 0:
+            role = "system"
+        else:
+            role = alternate_roles[(index - 1) % 2]
+
+        # Clean the message content by removing leading/trailing whitespace and separators
+        clean_content = message_content.strip()
+        for separator in separators:
+            clean_content = clean_content.replace(separator.strip("'"), '', 1).strip()
+
+        # Append the cleaned message with its role to the list
+        recovered_messages.append({"role": role, "content": clean_content})
+
+    return recovered_messages
+
+def recover_chat_messages(tokenized_chat, tokenizer):
+    """
+    Given a tokenized_chat string and a tokenizer, returns the list of message dictionaries.
+    """
+    jinja_template = tokenizer.chat_template
+    separators = extract_separators(jinja_template)
+    eos_token = eos_token = detect_eos_token(jinja_template, tokenizer)
+    recovered_messages = recover_messages(tokenized_chat, separators, eos_token)
+    return recovered_messages
+
+# Example usage
+if __name__ == "__main__":
+    checkpoint = "Qwen/Qwen1.5-0.5B"
+    tokenizer = AutoTokenizer.from_pretrained(checkpoint)
+    
+    messages = [
+        {
+            "role": "system",
+            "content": "You are a friendly chatbot who always responds in the style of a pirate",
+        },
+        {"role": "user", "content": "How many helicopters can a human eat in one sitting?"},
+    ]
+    tokenized_chat = tokenizer.apply_chat_template(messages, tokenize=False)
+    print(tokenized_chat)
+    
+    recovered_messages = recover_chat_messages(tokenized_chat, tokenizer)
+    print(recovered_messages)