Overview
The Kraken model and Architecture Kraken is a joint effort between Cognitive Computations, VAGO Solutions and Hyperspace.ai.
Created by Fernando Fernandes Neto, David Golchinfar, Lucas Atkins and Eric Hartford
The Kraken model combining the best Python, SQL, Function Calling, Reasoning and foreign Models so far.
The Kraken Architecture is a sophisticated machine learning framework designed for dynamic text generation tasks. It utilizes the Hugging Face transformers library to orchestrate multiple causal language models (CLMs) and intelligently route input through different models based on the context and content of the input text. The architecture is powered by a custom configuration class (KrakenConfig) that facilitates the integration and management of various components such as tokenizers, models, and routing mechanisms.
Features
Dynamic Model Routing: Uses a sequence classification model to route inputs to the most suitable language model based on the input's characteristics. Multiple Language Models: Supports integration of various pre-trained causal language models, allowing for flexible, context-appropriate responses. Customizable Templates: Includes support for input formatting using predefined templates, enhancing the model's adaptability to different conversational contexts. Extensible Configuration: Leverages a custom configuration setup that can be easily extended and adapted for various use cases involving causal language modeling.
Selected Models as Experts:
"Reasoning Expert": "microsoft/Phi-3-medium-128k-instruct",
"Function Calling Expert": "gorilla-llm/gorilla-openfunctions-v2",
"Python Expert": "ise-uiuc/Magicoder-S-DS-6.7B",
"SQL Expert": "defog/llama-3-sqlcoder-8b",
"German Expert": "VAGOsolutions/Llama-3-SauerkrautLM-8b-Instruct"
How to load and call Kraken model :
from transformers import AutoConfig, AutoModelForCausalLM
from configuration_kraken import KrakenConfig
from modeling_kraken import KrakenForCausalLM
AutoConfig.register("kraken", KrakenConfig)
AutoModelForCausalLM.register(KrakenConfig, KrakenForCausalLM)
device = "cuda:0" ## Setup "cuda:0" if NVIDIA, "mps" if on Mac
# Load the model and config:
config = AutoConfig.from_pretrained("./kraken_model")
model = AutoModelForCausalLM.from_pretrained("./kraken_model", config=config, trust_remote_code=True)
Call the Reasoning expert:
messages = [
{'role': 'system', 'content': '"You are a helpful AI Assistant'},
{'role': 'user', 'content': "Find the mass percentage of Ba in BaO"}
]
tokenizer = model.tokenizer
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
output_ids = model.generate(input_ids, max_length=250)
print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))
Call the Function Calling Expert:
messages = [
{'role': 'system', 'content': """You are a helpful assistant with access to the following functions. Use them if required -
{ "name": "calculate_area", "description": "Calculate the area of a shape", "parameters": { "type": "object", "properties": { "shape": { "type": "string", "description": "The type of shape (circle, rectangle, triangle, etc.)" }, "dimensions": { "type": "object", "properties": { "length": { "type": "number", "description": "The length of the shape" }, "width": { "type": "number", "description": "The width of the shape" }, "base": { "type": "number", "description": "The base of the shape" }, "height": { "type": "number", "description": "The height of the shape" } } } }, "required": [ "shape" ] } }"""},
{'role': 'user', 'content': """I need to calculate the area of a rectangle. The length is 5 and the width is 3."""}
]
tokenizer = model.tokenizer
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda:0")
output_ids = model.generate(input_ids ,temperature=0.1, do_sample=True, top_p=0.9,top_k=20, max_length=500)
print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))
Call the Python Expert:
messages = [
{'role': 'system', 'content': ''},
{'role': 'user', 'content': """Create a python function to calculate the sum of a sequence of integers.
[1, 2, 3, 4, 5]"""}
]
tokenizer = model.tokenizer
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
print(input_text)
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda:0")
output_ids = model.generate(input_ids ,temperature=0.6, do_sample=True, top_p=0.9,top_k=20, max_length=400)
print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))
Call the SQL expert:
messages = [
{'role': 'system', 'content': 'You are a helpul AI assistant.'},
{'role': 'user', 'content': """Generate a SQL query to answer this question: What is the total volume of timber sold by each salesperson, sorted by salesperson?
DDL statements:
CREATE TABLE salesperson (salesperson_id INT, name TEXT, region TEXT); INSERT INTO salesperson (salesperson_id, name, region) VALUES (1, 'John Doe', 'North'), (2, 'Jane Smith', 'South'); CREATE TABLE timber_sales (sales_id INT, salesperson_id INT, volume REAL, sale_date DATE); INSERT INTO timber_sales (sales_id, salesperson_id, volume, sale_date) VALUES (1, 1, 120, '2021-01-01'), (2, 1, 150, '2021-02-01'), (3, 2, 180, '2021-01-01');"""}
]
tokenizer = model.tokenizer
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to(device)
output_ids = model.generate(input_ids ,temperature=0.6, do_sample=True, top_p=0.9,top_k=20, max_length=500)
print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))
Call the German expert:
messages = [
{'role': 'system', 'content': 'Du bist ein freundlicher und hilfreicher deutscher KI-Assistent'},
{'role': 'user', 'content': "Ich hoffe es geht dir gut?"}
]
tokenizer = model.tokenizer
input_text = tokenizer.apply_chat_template(messages, tokenize=False)
input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda:0")
output_ids = model.generate(input_ids, max_length=150)
print(model.expert_tokenizer(text=input_text).decode(output_ids[0], skip_special_tokens=True))
Switch expert and or quantization:
Go into the config file of the kraken_model folder
"models": {
"expert1": "microsoft/Phi-3-medium-128k-instruct", # Switch to a Resoning Expert of your choice
"expert2": "gorilla-llm/gorilla-openfunctions-v2", # Switch to a Function Calling Expert of your choice
"expert3": "ise-uiuc/Magicoder-S-DS-6.7B", # Switch to a Python Expert of your choice
"expert4": "defog/llama-3-sqlcoder-8b", # Switch to a SQL Expert of your choice
"expert5": "VAGOsolutions/Llama-3-SauerkrautLM-8b-Instruct" # Switch to a German Expert of your choice
},
# Currently supported: "4bit","8bit" and "awq"
"quantization": {
"expert1": null,
"expert2": null,
"expert3": null,
"expert4": null,
"expert5": null
},
"router": "kraken_router",
# Adjust the tokenizer to your selected model
"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"
}
Cite As
Fernando Fernandes Neto, David Golchinfar, Lucas Atkins, Eric Hartford - Kraken: An OpenSource Collection of Experts Model, 2024