IUS-SWE-Agent-1.6b

Model Description

IUS-SWE-Agent-1.6b is an autonomous coding agent specialized in generating multi-step debugging trajectories (ReAct format: THOUGHT, ACTION, OBSERVATION).

It is built upon a novel Universal Transformer architecture featuring DeepSeek-style Mixture of Experts (MoE) and Multi-Head Latent Attention (MLA). By looping a single 132M parameter block 12 times, it achieves the computational depth of a 1.6 Billion parameter model while maintaining a tiny VRAM footprint.

Architecture Details

• Physical Parameters: ~132.2M
• Effective Depth: ~1.6B (12 recursive passes)
• Attention: Multi-Head Latent Attention (MLA) with KV compression
• Feed-Forward: DeepSeek MoE (1 Shared + 8 Routed Experts, Top-K=2)
• Context Window: 1024 tokens

Training Data

Trained on SWE-agent style trajectories (ISSUE -> THOUGHT -> ACTION -> OBSERVATION).

How to Use This Model

Because this is a custom architecture, you must define the model classes in your Python script before loading the weights.

1. Define the Architecture

Copy the architecture classes (RMSNorm, SwiGLU, MultiHeadLatentAttention, DeepSeekMoE, UniversalBlock, CodingAgent) from the training script.

2. Load and Generate

import torch, json, tiktoken
from huggingface_hub import hf_hub_download

repo_id = 'antontuzovAI/IUS-SWE-Agent-1.6b'
config_path = hf_hub_download(repo_id=repo_id, filename='config.json')
weights_path = hf_hub_download(repo_id=repo_id, filename='model_weights.pth')

with open(config_path, 'r') as f:
    config = json.load(f)

enc = tiktoken.get_encoding('gpt2')
model = CodingAgent(**config)
model.load_state_dict(torch.load(weights_path, map_location='cpu'))
model.eval()

def generate_agent_plan(issue, max_new_tokens=250):
    prompt = f'<AGENT>\nISSUE: {issue}\n'
    tokens = enc.encode(prompt, allowed_special={'<|endoftext|>'})
    x = torch.tensor([tokens])
    with torch.no_grad():
        for _ in range(max_new_tokens):
            logits = model(x)
            next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
            x = torch.cat([x, next_token], dim=1)
            if enc.decode([next_token.item()]) == '</AGENT>': break
    return enc.decode(x[0].tolist()).split('</AGENT>')[0].strip()

print(generate_agent_plan('Fix the KeyError in the login function.'))