Llada 8B Tools

🚀 LLaDA-8B-Tools

本倉庫包含一個基於 GSAI-ML/LLaDA-8B-Instruct 模型的變體，由 Proximile LLC 進行微調，以增強其工具調用能力。Proximile 專注於為中小型企業提供安全的本地部署 AI 解決方案。

📅 更新時間線

• 2025年5月14日 – 首次公開發布。訓練示例中缺少填充生成窗口其餘部分的填充標記。
• 2025年5月17日 – 修補訓練腳本以包含正確的填充；將更新後的模型權重推送到此倉庫。

✨ 關於LLaDA

LLaDA（Large Language Diffusion with mAsking）是一種新穎的語言模型架構，它使用離散擴散進行文本生成。與傳統的自迴歸模型不同，LLaDA 通過迭代去噪過程生成文本，根據置信度分數逐步用預測標記替換掩碼標記。

📚 模型描述

這個合併的 LoRA 模型經過訓練，以提高 LLaDA 處理工具調用任務的能力，包括：

• 為工具調用生成合適的 JSON
• 處理工具響應數據
• 根據工具輸出提供有用的答案

訓練詳情

• 基礎模型：GSAI-ML/LLaDA-8B-Instruct
• 訓練方法：使用 LoRA 的監督微調（SFT）
• LoRA 配置：
  • 秩（r）：128
  • Alpha：256
  • 目標模塊：q_proj、k_proj、v_proj、gate_proj
• 訓練數據：ToolACE 數據集的修改子集。

📦 安裝指南

pip install transformers peft torch bitsandbytes

💻 使用示例

基礎用法

from transformers import AutoTokenizer, AutoModel
from peft import PeftModel

# 加載基礎模型和分詞器
model_name = "Proximile/LLaDA-8B-Tools"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map="auto")

高級用法

以下是一個使用該模型進行帶有工具調用的聊天完成的完整示例：

import torch
import json
from transformers import AutoTokenizer, AutoModel

# 常量
MASK_TOKEN_ID = 126336

def add_gumbel_noise(logits, temperature):
    '''
    The Gumbel max is a method for sampling categorical distributions.
    For diffusion models, low-precision Gumbel Max affects generation quality.
    '''
    if temperature <= 0:
        return logits
        
    logits = logits.to(torch.float64)
    noise = torch.rand_like(logits, dtype=torch.float64)
    gumbel_noise = (- torch.log(noise)) ** temperature
    return logits.exp() / gumbel_noise

def get_num_transfer_tokens(mask_index, steps):
    '''
    In the reverse process, we precompute the number of tokens to transition at each step.
    '''
    mask_num = mask_index.sum(dim=1, keepdim=True)
    
    # Ensure we have at least one step
    if steps == 0:
        steps = 1
        
    base = mask_num // steps
    remainder = mask_num % steps
    
    num_transfer_tokens = torch.zeros(mask_num.size(0), steps, device=mask_index.device, dtype=torch.int64) + base
    
    for i in range(mask_num.size(0)):
        if remainder[i] > 0:
            num_transfer_tokens[i, :remainder[i]] += 1
            
    return num_transfer_tokens

def generate(model, prompt, steps=128, gen_length=128, block_length=32, temperature=0.,
             remasking='low_confidence', mask_id=MASK_TOKEN_ID):
    '''
    Generate text using LLaDA's diffusion-based generation process.
    '''
    device = next(model.parameters()).device
    prompt = prompt.to(device)
    
    x = torch.full((1, prompt.shape[1] + gen_length), mask_id, dtype=torch.long).to(device)
    x[:, :prompt.shape[1]] = prompt.clone()
    
    prompt_index = (x != mask_id)
    
    assert gen_length % block_length == 0
    num_blocks = gen_length // block_length
    
    assert steps % num_blocks == 0
    steps_per_block = steps // num_blocks
    
    for num_block in range(num_blocks):
        block_mask_index = (x[:, prompt.shape[1] + num_block * block_length: prompt.shape[1] + (num_block + 1) * block_length:] == mask_id)
        num_transfer_tokens = get_num_transfer_tokens(block_mask_index, steps_per_block)
        
        for i in range(steps_per_block):
            mask_index = (x == mask_id)
            if not mask_index.any():
                break
                
            outputs = model(x)
            logits = outputs.logits
            
            logits_with_noise = add_gumbel_noise(logits, temperature=temperature)
            x0 = torch.argmax(logits_with_noise, dim=-1)  # b, l
            
            if remasking == 'low_confidence':
                p = torch.nn.functional.softmax(logits.to(torch.float64), dim=-1)
                x0_p = torch.squeeze(
                    torch.gather(p, dim=-1, index=torch.unsqueeze(x0, -1)), -1)  # b, l
            elif remasking == 'random':
                x0_p = torch.rand((x0.shape[0], x0.shape[1]), device=x0.device)
            else:
                raise NotImplementedError(remasking)
            
            x0_p[:, prompt.shape[1] + (num_block + 1) * block_length:] = -float('inf')
            
            x0 = torch.where(mask_index, x0, x)
            confidence = torch.where(mask_index, x0_p, -float('inf'))
            
            transfer_index = torch.zeros_like(x0, dtype=torch.bool, device=x0.device)
            for j in range(confidence.shape[0]):
                _, select_index = torch.topk(confidence[j], k=num_transfer_tokens[j, i])
                transfer_index[j, select_index] = True
            x[transfer_index] = x0[transfer_index]
    
    return x

def chat_completion(model, tokenizer, messages, temperature=0.1, gen_length=128, steps=128):
    """
    Generate a chat completion.
    
    Args:
        model: The LLaDA tool calling model
        tokenizer: The tokenizer
        messages: List of message dictionaries with 'role' and 'content' keys
        temperature: Temperature for generation (0 for greedy)
        gen_length: Maximum length of generated text
        steps: Number of denoising steps
        
    Returns:
        The generated response text
    """
    # Format input for the model
    formatted_input = tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=True
    )
    
    # Tokenize input
    input_ids = tokenizer(formatted_input, return_tensors="pt")["input_ids"]
    
    # Generate response
    with torch.no_grad():
        output_ids = generate(
            model, 
            input_ids, 
            steps=steps,
            gen_length=gen_length,
            block_length=32,
            temperature=temperature,
            remasking='low_confidence'
        )
    
    # Decode the generated output
    generated_text = tokenizer.decode(output_ids[0, input_ids.shape[1]:], skip_special_tokens=False).split("<|")[0]
    return generated_text

# Example usage
if __name__ == "__main__":
    # Load the base model and tokenizer
    model_name = "Proximile/LLaDA-8B-Tools"
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    model = AutoModel.from_pretrained(model_name, trust_remote_code=True, device_map="auto")
    
    # Define tool calling function schema
    tool_schema = [
        {
            "type": "function",
            "function": {
                "name": "get_weather",
                "description": "Get the current weather in a given location",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "location": {
                            "type": "string",
                            "description": "The city and state, e.g. San Francisco, CA"
                        },
                        "unit": {
                            "type": "string",
                            "enum": ["celsius", "fahrenheit"],
                            "description": "The unit of temperature"
                        }
                    },
                    "required": ["location", "unit"]
                }
            }
        }
    ]
    
    # Create conversation with system prompt including tool description
    system_prompt = """You are a helpful assistant with tool calling capabilities. When you receive a tool call response, use the output to format an answer to the orginal user question.

If you choose to use one or more of the following tool functions, respond with a list of JSON function calls, each with the proper arguments that best answers the given prompt.

Each tool request within the list should be in the exact format {"name": function name, "parameters": {dictionary of argument names and values}}. Do not use variables. Just a list of two-key dictionaries, each starting with the function name, followed by a dictionary of parameters.

Here are the tool functions available to you:

""" + json.dumps(tool_schema, indent=4) + """

After receiving the results back from a function call, you have to formulate your response to the user. If the information needed is not found in the returned data, either attempt a new function call, or inform the user that you cannot answer based on your available knowledge. The user cannot see the function results. You have to interpret the data and provide a response based on it.

If the user request does not necessitate a function call, simply respond to the user's query directly."""
    
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": "What's the weather like in New York?"}
    ]
    
    # Generate assistant response (expecting tool call)
    assistant_response = chat_completion(model, tokenizer, messages)
    print(f"Assistant: {assistant_response}")
    
    # Mock tool response
    tool_response = json.dumps({
        "location": "New York, NY",
        "temperature": 72,
        "unit": "fahrenheit",
        "condition": "Partly Cloudy",
        "humidity": 65,
        "wind_speed": 8,
        "wind_direction": "NE"
    })
    
    # Add assistant and tool responses to the conversation
    messages.append({"role": "assistant", "content": assistant_response})
    messages.append({"role": "ipython", "content": tool_response})
    
    # Generate final assistant response
    final_response = chat_completion(model, tokenizer, messages)
    print(f"Assistant (with tool data): {final_response}")

# Assistant: [{"name": "get_weather", "parameters": {"location": "New York", "unit": "fahrenheit"}}]
# Assistant (with tool data): The current weather in New York is as follows:
# - Temperature: 72°F
# - Weather Condition: Partly Cloudy
# - Humidity: 65%
# - Wind Speed: 8 miles per hour
# - Wind Direction: Northeast

⚠️ 侷限性

⚠️ 重要提示

• LLaDA 基於擴散的生成方式與標準大語言模型不同，在某些上下文中可能表現不同。
• 模型仍可能產生幻覺或生成不正確的工具調用格式。
• 工具調用的格式必須與示例中所示的格式（這是 官方 llama 3.1 格式 的修改版本）完全匹配。

📖 引用

如果您在研究中使用此模型，請同時引用原始的 LLaDA 論文和此適配器：

@misc{llada-8b-tools,
  author = {Proximile LLC},
  title = {LLaDA-8B-Tools},
  year = {2025},
  publisher = {Hugging Face},
  howpublished = {\url{https://huggingface.co/Proximile/LLaDA-8B-Tools}}
}

🏢 關於 Proximile LLC

Proximile LLC 為中小型企業提供安全、經濟高效且私密的 AI 解決方案。我們專注於：

• 本地部署 AI 推理 解決方案，確保無與倫比的隱私性。
• 經濟高效的硬件配置，包括 Jetson Orin Nano Super。
• 安全的本地 AI 應用程序，包括聊天機器人、RAG 系統和自定義 AI 工具。
• 專業服務，用於合規與治理、知識管理和 IT 自動化。

訪問 proximile.llc 以瞭解更多關於我們為您的企業提供的安全本地 AI 解決方案。

📄 許可證

此適配器與基礎 LLaDA 模型使用相同的許可證發佈。