lv12/k3-multimodal-6k

K3 Multimodal Language Model

A 149M parameter implementation of the Kimi K3 architecture with support for text, images, and audio. Trained for 6000 steps on multimodal data.

Installation

pip install torch transformers pillow

Usage

Text Generation

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

model = AutoModelForCausalLM.from_pretrained(
    "lv12/k3-multimodal-6k",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained("lv12/k3-multimodal-6k")

device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

prompt = "The future of AI is"
inputs = tokenizer(prompt, return_tensors="pt").to(device)

outputs = model.generate(**inputs, max_new_tokens=50, temperature=0.7)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

With Images

from PIL import Image
import torchvision.transforms as transforms

image = Image.open("example.jpg").convert("RGB")
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image_tensor = transform(image).unsqueeze(0).unsqueeze(0).to(device)

outputs = model(
    input_ids=inputs.input_ids,
    images=image_tensor,
    has_visual=torch.ones(1, dtype=torch.bool, device=device)
)

Training

from transformers import Trainer, TrainingArguments
from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train[:1000]")
tokenized = dataset.map(
    lambda x: tokenizer(x["text"], truncation=True, max_length=512, padding="max_length"),
    batched=True, remove_columns=["text"]
)

trainer = Trainer(
    model=model,
    args=TrainingArguments(
        output_dir="./output",
        num_train_epochs=3,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        learning_rate=6e-4,
        warmup_ratio=0.03,
        save_steps=500,
    ),
    train_dataset=tokenized,
)

trainer.train()

Model Details

• 149M parameters (sparse MoE: 256 experts, 4 active per token)
• 163,840 token vocabulary (tiktoken BPE)
• 1024 token context length
• Multimodal: text, images (ViT), audio (Whisper-style)

Citation

@article{liu2024k3,
  title={K3: A New Era of Multimodal Large Language Models},
  author={Liu, Yang and others},
  journal={arXiv preprint arXiv:2407.24653},
  year={2024}
}