Yi

Quick start

💡 Tip: If you want to get started with the Yi model and explore different methods for inference, check out the Yi Cookbook.

Quick start - pip

This tutorial guides you through every step of running Yi-34B-Chat locally on an A800 (80G) and then performing inference.

Step 0: Prerequisites

• Make sure Python 3.10 or a later version is installed.

• If you want to run other Yi models, see software and hardware requirements.

Step 1: Prepare your environment

To set up the environment and install the required packages, execute the following command.

git clone https://github.com/01-ai/Yi.git
cd yi
pip install -r requirements.txt

Step 2: Download the Yi model

You can download the weights and tokenizer of Yi models from the following sources:

• Hugging Face
• ModelScope
• WiseModel

Step 3: Perform inference

You can perform inference with Yi chat or base models as below.

Perform inference with Yi chat model

1. Create a file named quick_start.py and copy the following content to it.

   from transformers import AutoModelForCausalLM, AutoTokenizer
   
   model_path = '<your-model-path>'
   
   tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)
   
   # Since transformers 4.35.0, the GPT-Q/AWQ model can be loaded using AutoModelForCausalLM.
   model = AutoModelForCausalLM.from_pretrained(
       model_path,
       device_map="auto",
       torch_dtype='auto'
   ).eval()
   
   # Prompt content: "hi"
   messages = [
       {"role": "user", "content": "hi"}
   ]
   
   input_ids = tokenizer.apply_chat_template(conversation=messages, tokenize=True, add_generation_prompt=True, return_tensors='pt')
   output_ids = model.generate(input_ids.to('cuda'))
   response = tokenizer.decode(output_ids[0][input_ids.shape[1]:], skip_special_tokens=True)
   
   # Model response: "Hello! How can I assist you today?"
   print(response)
   

2. Run quick_start.py.

   python quick_start.py
   

   Then you can see an output similar to the one below. 🥳

   Hello! How can I assist you today?
   

Perform inference with Yi base model

• Yi-34B

  The steps are similar to pip - Perform inference with Yi chat model.

  You can use the existing file text_generation.py.

  python demo/text_generation.py  --model <your-model-path>
  

  Then you can see an output similar to the one below. 🥳

  Output. ⬇️

  
  

  Prompt: Let me tell you an interesting story about cat Tom and mouse Jerry,

  Generation: Let me tell you an interesting story about cat Tom and mouse Jerry, which happened in my childhood. My father had a big house with two cats living inside it to kill mice. One day when I was playing at home alone, I found one of the tomcats lying on his back near our kitchen door, looking very much like he wanted something from us but couldn’t get up because there were too many people around him! He kept trying for several minutes before finally giving up...

• Yi-9B

  Input

  from transformers import AutoModelForCausalLM, AutoTokenizer
  
  MODEL_DIR = "01-ai/Yi-9B"
  model = AutoModelForCausalLM.from_pretrained(MODEL_DIR, torch_dtype="auto")
  tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, use_fast=False)
  
  input_text = "# write the quick sort algorithm"
  inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
  outputs = model.generate(**inputs, max_length=256)
  print(tokenizer.decode(outputs[0], skip_special_tokens=True))
  

  Output

  # write the quick sort algorithm
  def quick_sort(arr):
      if len(a