许可证：Apache-2.0
标签：

• jamba
  数据集：
• Locutusque/hercules-v4.0
  基础模型：ai21labs/Jamba-v0.1
  管道标签：文本生成

Jamba-Hercules

使用数据集：Locutusque/hercules-v4.0

- 前1万条示例

示例输出：

<|startoftext|><|im_start|>system  
你是一个乐于助人的AI助手。  
<|im_end|>  
<|im_start|>user  
一棵树需要离开地球，成为太空生物。它该如何实现这个目标？  
<|im_end|>  
<|im_start|>assistant  
让一棵树成为太空生命是一个极其复杂的过程，需要大量的准备、协调和资金支持。第一步是评估这棵树是否适合太空旅行，包括对其大小、健康状况和适应性的全面检查。  
一旦被选为候选，接下来需要为太空飞行做准备。这包括训练树木适应低重力环境、学习在零重力下移动，并确保其在太空中获得足够的营养。适应过程可能需要数月甚至数年，具体取决于树种和任务需求。  
准备就绪后，树木需要通过火箭或其他方式进入太空。这段旅程漫长且充满风险。  
到达太空后，树木需要宇航员的密切照料，包括提供养分、保护其免受辐射和极端温度的影响。  
随着时间的推移，树木可能会在形态和行为上发生变化，例如长出新型枝叶或利用根部探索新环境。  
这一过程中有许多挑战需要克服。<|im_end|>  
<|endoftext|>  

推理

※ - 此代码会自动为输入添加"<|startoftext|>"特殊标记。测试表明，缺少该标记会导致输出错误。

!pip install -qqq transformers>=4.39.0 mamba-ssm causal-conv1d>=1.2.0 accelerate bitsandbytes --progress-bar off  
!pip install flash-attn --no-build-isolation  

import torch  
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig  

double_quant_config = BitsAndBytesConfig(  
   load_in_4bit=True,  
   bnb_4bit_use_double_quant=True,  
   bnb_4bit_compute_dtype=torch.float16  
)  

model = AutoModelForCausalLM.from_pretrained(  
    "Severian/Jamba-Hercules",  
    device_map="auto",  
    trust_remote_code=True,  
    torch_dtype=torch.bfloat16,  
    attn_implementation="flash_attention_2",  
    quantization_config=double_quant_config,  
)  
tokenizer = AutoTokenizer.from_pretrained("Severian/Jamba-Hercules")  

input_text = """<|im_start|>system  
你是一个乐于助人的AI助手。  
<|im_end|>  
<|im_start|>user  
一棵树需要离开地球，成为太空生物。它该如何实现这个目标？  
<|im_end|>  
<|im_start|>assistant  
"""  

input_ids = tokenizer(input_text, return_tensors='pt').to(model.device)["input_ids"]  

outputs = model.generate(input_ids, max_new_tokens=1024, temperature=0.0, repetition_penalty=1.1)  

print(tokenizer.batch_decode(outputs)[0])  

训练

Hercules-v4.0:

首次测试:

• 1250步（A100显卡5小时）
• 最终损失值：0.98

超参数

lora_config = LoraConfig(  
    r=16,  
    lora_alpha=32,  
    target_modules=["embed_tokens", "x_proj", "in_proj", "out_proj"],  
    lora_dropout=0.05,  
    task_type="CAUSAL_LM",  
    bias="none"  
)  

trainer = SFTTrainer(  
    model=model,  
    train_dataset=train_dataset,  
    dataset_text_field="text",  
    max_seq_length=max_seq_length,  
    tokenizer=tokenizer,  
    args=TrainingArguments(  
        num_train_epochs=1,  
        lr_scheduler_type='cosine',  
        learning_rate=0.0002,  
        per_device_train_batch_size=1,  
        gradient_accumulation_steps=8,  
        gradient_checkpointing=True,  
        warmup_steps=10,  
        weight_decay=0.01,  
        fp16=not torch.cuda.is_bf16_supported(),  
        bf16=torch.cuda.is_bf16_supported(),  
        logging_steps=1,  
        save_steps=200,  
        output_dir="outputs",  
        optim="adamw_bnb_8bit",  
        adam_epsilon=0.00001,  
        adam_beta2=0.95,  
        max_grad_norm=1.0,  
        seed=42,  
    ),  
)