SurjoLabs

Surjo-50M-SFT-Only

Surjo-50M-SFT-Only is the instruction-tuned checkpoint of Surjo-50M, fine-tuned for basic dialogue and instruction-following tasks.

Details

Usage

Installation

pip install torch transformers accelerate
pip install -U git+https://github.com/fla-org/flash-linear-attention

Chat Inference

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "SurjoLabs/Surjo-50m-SFT-Only"

tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
    device_map="auto",
    trust_remote_code=True,
)

messages = [
    {"role": "user", "content": "Explain what a neural network is in two sentences."}
]

inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
).to(model.device)

with torch.no_grad():
    outputs = model.generate(
        inputs,
        max_new_tokens=128,
        do_sample=True,
        temperature=0.7,
        top_p=0.9,
    )

response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
print(response)

License

Licensed under the Apache 2.0 License.