abhinand / sarvam-105b-bf16

See footnote for evaluation details.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig

model_name = "sarvamai/sarvam-105b"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, device_map="auto")

def generate_text(
    prompt: str,
    max_new_tokens: int = 2048,
    temperature: float = 0.8,
    top_p: float = 0.95,
    repetition_penalty: float = 1.0,
) -> None:
    inputs = tokenizer(prompt, return_tensors="pt").to("cuda:0")

    generation_config = GenerationConfig(
        max_new_tokens=max_new_tokens,
        repetition_penalty=repetition_penalty,
        temperature=temperature,
        top_p=top_p,
        do_sample=True,
    )

    with torch.no_grad():
        output_ids = model.generate(
            input_ids=inputs["input_ids"],
            attention_mask=inputs["attention_mask"],
            generation_config=generation_config,
        )
    return tokenizer.decode(output_ids[0], skip_special_tokens=True)

prompts = [
    "Which country won the FIFA World Cup in 2012?",
]

for prompt in prompts:
    templated_prompt = tokenizer.apply_chat_template(
      [{"role": "user", "content": prompt}],
      tokenize=False,
      add_generation_prompt=True,
      enable_thinking=True
    )
    output = generate_text(templated_prompt, max_new_tokens=512)
    print("Prompt: ", prompt)
    print("Generated text: ", output)
    print("=" * 100)

Install latest SGLang from source

git clone https://github.com/sgl-project/sglang.git
cd sglang
pip install -e "python[all]"

Instantiate model and Run

import sglang as sgl
from transformers import AutoTokenizer

model_path = "abhinand/sarvam-105b-bf16"
engine = sgl.Engine(
    model_path=model_path,
    tp_size=4,
    mem_fraction_static=0.70,
    trust_remote_code=True,
    dtype="bfloat16",
    moe_runner_backend="flashinfer_cutedsl",
    prefill_attention_backend="fa3",
    decode_attention_backend="flashmla",
    disable_radix_cache=False,
)

sampling_params = {
    "temperature": 0.8,
    "max_new_tokens": 2048,
    "repetition_penalty": 1.0,
}

prompts = [
    "Which band released the album Dark Side of the Moon in 1973?",
]

outputs = engine.generate([
    tokenizer.apply_chat_template([
        {"role": "user", "content": prompt}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True)
        for prompt in prompts], 
    sampling_params)
for p, o in zip(prompts, outputs):
    print("Prompt: ", p)
    print("Generated text: ", o['text'])
    print("=" * 100)

Note: currently a PR is open for native support for the Sarvam models in vLLM ( link ). Therefore, we have 2 options here.

Option 1: install from source (hard)

• Use the custom fork here: link
• Follow the instructions here to install from source: link

Option 2: hot-patch (easy)

• Run hotpatch_vllm.py
• This will do the following:
  • install vllm=0.15.0
  • add 2 model entries to registry.py
  • download the model executors for sarvam-105b and sarvam-30b

Once this is done, you can run vLLM as usual

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

model_path = "abhinand/sarvam-105b-bf16"
tokenizer = AutoTokenizer.from_pretrained(model_path)
llm = LLM(model=model_path, 
            trust_remote_code=True, 
            max_model_len=2048, 
            tensor_parallel_size=8, 
            max_num_seqs=16,
        )
sampling_params = SamplingParams(
                    temperature=0.8, 
                    max_tokens=2048, 
                    repetition_penalty=1.0,
                    spaces_between_special_tokens=True
                )

prompts = [
    "Which artist painted The Persistence of Memory (the melting clocks)?",
]

outputs = llm.generate([
    tokenizer.apply_chat_template([
        {"role": "user", "content": prompt}],
        tokenize=False,
        add_generation_prompt=True,
        enable_thinking=True)
        for prompt in prompts], 
    sampling_params)
for p, o in zip(prompts, outputs):
    print("Prompt: ", p)
    print("Generated text: ", o.outputs[0].text)
    print("=" * 100)