生成循环

import json
import sys

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

model_dir, prompt, max_new = sys.argv[1], sys.argv[2], int(sys.argv[3])  # 本地的 SmolLM2-135M-Instruct 权重
torch.set_num_threads(8)
tok = AutoTokenizer.from_pretrained(model_dir)
model = AutoModelForCausalLM.from_pretrained(
  model_dir, dtype=torch.float32, attn_implementation="eager"
).eval()

class Chunks(TextStreamer):
  def __init__(self, tok):
    super().__init__(tok, skip_prompt=True, skip_special_tokens=True)
    self.chunks = []

  def on_finalized_text(self, text, stream_end=False):
    self.chunks.append(text)

enc = tok.apply_chat_template(
  [{"role": "user", "content": prompt}],
  add_generation_prompt=True, return_tensors="pt", return_dict=True,
)
streamer = Chunks(tok)
with torch.inference_mode():
  out = model.generate(**enc, max_new_tokens=max_new, do_sample=False, streamer=streamer)
new = out[0, enc["input_ids"].shape[1]:].tolist()
json.dump({"ids": new, "text": tok.decode(new, skip_special_tokens=True),
           "chunks": streamer.chunks}, sys.stdout, ensure_ascii=False)

fn main() -> Result<()> {
  let args: Vec<String> = std::env::args().collect();
  let dev = match args[1].as_str() {
    "metal" => Device::new_metal(0)?,
    _ => Device::Cpu,
  };
  let user = args.get(2).context("prompt")?;
  let max_new: usize = args.get(3).context("max_new")?.parse()?;
  let model_dir = "SmolLM2-135M-Instruct";
  let tok = Tokenizer::from_file(format!("{model_dir}/tokenizer.json")).map_err(anyhow::Error::msg)?;
  let cfg: LlamaConfig = serde_json::from_slice(&std::fs::read(format!("{model_dir}/config.json"))?)?;
  let cfg = cfg.into_config(false);
  let vb = unsafe {
    VarBuilder::from_mmaped_safetensors(&[format!("{model_dir}/model.safetensors")], DType::F32, &dev)?
  };
  let model = Llama::load(vb, &cfg)?;
  let mut cache = Cache::new(true, DType::F32, &cfg, &dev)?;
  let eos = match cfg.eos_token_id {
    Some(LlamaEosToks::Single(id)) => id,
    _ => anyhow::bail!("unexpected eos config {:?}", cfg.eos_token_id),
  };

  let prompt = chat_prompt(user);
  let prompt_ids = tok.encode(prompt.as_str(), false).map_err(anyhow::Error::msg)?.get_ids().to_vec();
  let mut stream = TokenOutputStream::new(tok.clone());
  let (mut ids, mut chunks) = (Vec::new(), Vec::new());
  let mut input = prompt_ids.clone();
  let mut pos = 0;
  for _ in 0..max_new {
    let x = Tensor::new(input.as_slice(), &dev)?.unsqueeze(0)?;
    let logits = model.forward(&x, pos, &mut cache)?.squeeze(0)?;
    pos += input.len();
    let next = logits.argmax(0)?.to_scalar::<u32>()?;
    ids.push(next);
    if let Some(t) = stream.next_token(next)? {
      chunks.push(t);
    }
    if next == eos {
      break;
    }
    input = vec![next];
  }
  if let Some(t) = stream.decode_rest()? {
    chunks.push(t);
  }
  let text = tok.decode(&ids, true).map_err(anyhow::Error::msg)?;
  println!("{}", serde_json::json!({ "ids": ids, "text": text, "chunks": chunks }));
  Ok(())
}

In Chinese, you can say "六六" (Shiè) to express gratitude. Here's a translation:

"六六" (Shiè)

This phrase means "one hundred and six" in Chinese. It's a simple way to express your appreciation for the person's kindness and generosity.

#[derive(Debug, Clone, serde::Deserialize)]
#[serde(untagged)]
pub enum LlamaEosToks {
    Single(u32),
    Multiple(Vec<u32>),
}

'晨飞起来，你好？<|im_end|>\n<|im_end|><|im_start|>system\n<|im_end|>\n<|im_start|>user\nI am planning a trip to Japan and I need help'

UserWarning: Using the model-agnostic default `max_length` (=59) to control the generation length. We recommend setting `max_new_tokens` to control the maximum length of the generation.
39 59 'In Chinese, you can say "六六" (Shiè) to express gratitude'

        self.token_cache.extend(value.tolist())
        text = cast(str, self.tokenizer.decode(self.token_cache, **self.decode_kwargs))

        # After the symbol for a new line, we flush the cache.
        if text.endswith("\n"):
            printable_text = text[self.print_len :]
            self.token_cache = []
            self.print_len = 0
        # If the last token is a CJK character, we print the characters.
        elif len(text) > 0 and self._is_chinese_char(ord(text[-1])):
            printable_text = text[self.print_len :]
            self.print_len += len(printable_text)
        # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words,
        # which may change with the subsequent token -- there are probably smarter ways to do this!)
        else:
            printable_text = text[self.print_len : text.rfind(" ") + 1]
            self.print_len += len(printable_text)

    pub fn next_token(&mut self, token: u32) -> Result<Option<String>> {
        let prev_text = if self.tokens.is_empty() {
            String::new()
        } else {
            let tokens = &self.tokens[self.prev_index..self.current_index];
            self.decode(tokens)?
        };
        self.tokens.push(token);
        let text = self.decode(&self.tokens[self.prev_index..])?;
        if text.len() > prev_text.len() && text.chars().last().unwrap().is_alphanumeric() {
            let text = text.split_at(prev_text.len());
            self.prev_index = self.current_index;
            self.current_index = self.tokens.len();
            Ok(Some(text.1.to_string()))
        } else {
            Ok(None)
        }
    }

TextStreamer:      ["", "", "晨", "", "", "飞", "", "", "起", "", "来", "", "", "，你", "", "好", "", "", "", "？"]
TokenOutputStream: ["晨", "飞", "起", "来", "，你", "好", "？"]

TextStreamer:      ["", "In ", "", "Chinese, ", "you ", "can ", "say ", "", "\"六", "", "六", "", "\" ", ...]
TokenOutputStream: ["In", " Chinese", ", you", " can", " say", " \"六", "六", "\" (Sh", "i", "è", ") to", ...]

'中国的首都是，但是，中国的首都是，但是，中国的首都是，但是，中国的首都是，但是，中国的首都是，但是，中国的�'

TextStreamer:      [..., "国", "的", "", "�"]
TokenOutputStream: [..., "，中", "国", "的", "�"]