Use transfomers tokenizer and streamer for python api (#388)

Author: zhenwei-intel

README.md

@@ -56,20 +56,36 @@ Below are the sample code to enable weight-only low precision inference. See mor
 
 ### INT4 Inference 
 ```python
+from transformers import AutoTokenizer
 from intel_extension_for_transformers.transformers import AutoModel, WeightOnlyQuantConfig
+
+model_name = "EleutherAI/gpt-j-6B"
+config = WeightOnlyQuantConfig(compute_dtype="int8", weight_dtype="int4")
 prompt = "Once upon a time, a little girl"
-config = WeightOnlyQuantConfig(compute_dtype="int8")
-model = AutoModel.from_pretrained("Intel/neural-chat-7b-v1-1", quantization_config=config)
-print(model.generate(prompt, max_new_tokens=30))
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+inputs = tokenizer(prompt, return_tensors="pt").input_ids
+
+model = AutoModel.from_pretrained(model_name, quantization_config=config)
+gen_tokens = model.generate(inputs, max_new_tokens=300)
+gen_text = tokenizer.batch_decode(gen_tokens)
 ```
 
 ### INT8 Inference
 ```python
+from transformers import AutoTokenizer
 from intel_extension_for_transformers.transformers import AutoModel, WeightOnlyQuantConfig
-prompt = "Once upon a time, a little girl"
+
+model_name = "EleutherAI/gpt-j-6B" 
 config = WeightOnlyQuantConfig(compute_dtype="bf16", weight_dtype="int8")
-model = AutoModel.from_pretrained("Intel/neural-chat-7b-v1-1", quantization_config=config)
-print(model.generate(prompt, max_new_tokens=30))
+prompt = "Once upon a time, a little girl"
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+inputs = tokenizer(prompt, return_tensors="pt").input_ids
+
+model = AutoModel.from_pretrained(model_name, quantization_config=config)
+gen_tokens = model.generate(inputs, max_new_tokens=300)
+gen_text = tokenizer.batch_decode(gen_tokens)
 ```
 
 ## 🎯Validated  Models

intel_extension_for_transformers/llm/runtime/graph/README.md

@@ -65,12 +65,19 @@ cmake --build . -j
 
 You can use Python API to run Hugging Face model simply. Here is the sample code:
 ```python
+from transformers import AutoTokenizer, TextStreamer
 from intel_extension_for_transformers.transformers import AutoModel, WeightOnlyQuantConfig
 model_name = "Intel/neural-chat-7b-v1-1"     # Hugging Face model_id or local model
 woq_config = WeightOnlyQuantConfig(compute_dtype="int8", weight_dtype="int4")
-model = AutoModel.from_pretrained(model_name, quantization_config=woq_config)
 prompt = "Once upon a time, a little girl"
-output = model.generate(prompt, max_new_tokens=30)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+inputs = tokenizer(prompt, return_tensors="pt").input_ids
+streamer = TextStreamer(tokenizer)
+
+model = AutoModel.from_pretrained(model_name, quantization_config=woq_config, trust_remote_code=True)
+outputs = model.generate(inputs, streamer=streamer, max_new_tokens=300)
+
 ```
 
 ### 3. Run LLM with Python Script

intel_extension_for_transformers/llm/runtime/graph/__init__.py

@@ -17,7 +17,7 @@
 import os
 from transformers import AutoConfig
 from intel_extension_for_transformers.llm.runtime.graph.scripts.convert import convert_model
-
+import torch
 model_maps = {"gpt_neox": "gptneox", "RefinedWebModel": "falcon"}
 
 class Model:
@@ -66,10 +66,12 @@ def init(self, model_name, **kwargs):
         # 1. convert model
         fp32_bin = "ne_{}_f32.bin".format(model_type)
         convert_model(model_name, fp32_bin, "f32")
+        assert(os.path.exists(fp32_bin), "Fail to convert pytorch model")
 
         # 2. quant model
         quant_bin = "ne_{}_q.bin".format(model_type)
         self.module.Model.quant_model(model_path = fp32_bin, out_path = quant_bin, **kwargs)
+        assert(os.path.exists(quant_bin), "Fail to quantize model")
 
         self.model_type = model_type
         self.bin_file = quant_bin
@@ -88,13 +90,22 @@ def quant_model(self, model_name, model_path, out_path, **kwargs):
         self.module.Model.quant_model(model_path = model_path,
                                     out_path = out_path, **kwargs)
 
-    def generate(self, prompt, streamer = None, sentence_mode = True, **kwargs):
-        # TODO support streamer
+    def generate(self, input_ids, streamer = None, **kwargs):
         if self.model is None:
             self.init_from_bin(self.model_type, self.bin_file, **kwargs)
-        
-        out = self.model.generate(prompt = prompt, sentence_mode = sentence_mode)
-        return out
+        # TODO support multi batch
+        assert(input_ids.shape[0] == 1, "Unsupport multi-batch input ids.")
+        if streamer:
+            ret = input_ids.tolist()
+            while not self.is_token_end():
+                out = self.model.generate(input_ids = input_ids.tolist()[0])
+                streamer.put(torch.tensor([out]))
+                ret[0].extend(out)
+            return ret
+        else:
+            ret = input_ids.tolist()
+            ret[0].extend(self.model.generate_tokens(input_ids = input_ids.tolist()[0]))
+            return ret
 
     def is_token_end(self):
         return self.model.is_token_end()

intel_extension_for_transformers/llm/runtime/graph/application/main_pybind.cpp

@@ -29,6 +29,7 @@
 #include <unordered_map>
 #include <utility>
 #include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
 #include "common.h"
 #include "models/model_utils/model_types.h"
 #include "models/model_utils/model_config.h"
@@ -57,7 +58,8 @@ class Model {
   void init_model(const std::string& model_path, int n_predict, int batch_size, int ctx_size, int seed, int threads,
                   float repeat_penalty, const std::string& post_process);
   void reinit();
-  std::string generate(const std::string& prompt, bool sentence_mode = true);
+  std::vector<int> generate(const std::vector<int>& input_ids);
+  std::vector<int> generate_tokens(const std::vector<int>& input_ids);
   bool is_token_end() { return token_eos; }
   static int quant_model(const std::string& model_path, const std::string& out_path, const std::string& weight_dtype,
                          const std::string& alg, int group_size, const std::string& scale_dtype,
@@ -73,8 +75,6 @@ class Model {
   std::vector<model_token> last_n_tokens;
   bool token_eos = false;
 
-  std::string generate_one_token(const std::string& prompt);
-  std::string generate_tokens(const std::string& prompt);
   int post_process(float* logits);
 };
 
@@ -108,10 +108,9 @@ void Model::reinit() {
   curr_input_ids.clear();
 }
 
-std::string Model::generate_one_token(const std::string& prompt) {
+std::vector<int> Model::generate(const std::vector<int>& input_ids) {
   if (curr_input_ids.empty()) {
-    auto embd_inp = ::model_tokenize(ctx, prompt, false);
-    curr_input_ids = embd_inp;
+    curr_input_ids = input_ids;
   }
   for (auto item : curr_input_ids) {
     last_n_tokens.erase(last_n_tokens.begin());
@@ -124,7 +123,7 @@ std::string Model::generate_one_token(const std::string& prompt) {
   int next_token_id = post_process(logits);
   curr_input_ids = {next_token_id};
 
-  if (next_token_id == ctx->vocab.eos_token_id || n_past - prompt.size() == params.n_predict) {
+  if (next_token_id == ctx->vocab.eos_token_id || n_past - input_ids.size() == params.n_predict) {
     token_eos = true;
   }
 
@@ -133,20 +132,17 @@ std::string Model::generate_one_token(const std::string& prompt) {
     token_eos = true;
   }
 
-  return next_token;
+  return {next_token_id};
 }
 
-std::string Model::generate_tokens(const std::string& prompt) {
-  int n_past = 0;
+std::vector<int> Model::generate_tokens(const std::vector<int>& input_ids) {
   int n_remain = params.n_predict;
-  int max_length = 512;
-  auto embd_inp = ::model_tokenize(ctx, prompt, false);
-  int n_eval = embd_inp.size();
-  std::vector<int> curr_input_ids(embd_inp);
   std::vector<int> output_ids;
-  output_ids.reserve(max_length);
-  std::string ret;
-  ret += prompt;
+
+  if (curr_input_ids.empty()) {
+    curr_input_ids = input_ids;
+  }
+
   while (output_ids.size() < n_remain) {
     for (auto item : curr_input_ids) {
       last_n_tokens.erase(last_n_tokens.begin());
@@ -158,24 +154,14 @@ std::string Model::generate_tokens(const std::string& prompt) {
     float* logits = model_get_logits(ctx);
     int next_token_id = post_process(logits);
     curr_input_ids = {next_token_id};
-
     output_ids.push_back(next_token_id);
-    ret += model_token_to_str(ctx, next_token_id);
-
-    if (next_token_id == model_token_eos()) {
+    if (next_token_id == ctx->vocab.eos_token_id || n_past - input_ids.size() == params.n_predict) {
+      token_eos = true;
       break;
     }
   }
 
-  return ret;
-}
-
-std::string Model::generate(const std::string& prompt, bool sentence_mode) {
-  if (sentence_mode) {
-    return generate_tokens(prompt);
-  }
-
-  return generate_one_token(prompt);
+  return output_ids;
 }
 
 int Model::post_process(float* logits) {
@@ -300,8 +286,8 @@ PYBIND11_MODULE(chatglm_cpp, m)
       .def("init_model", &Model::init_model, "initial model with model path and parameters", py::arg("model_path"),
            py::arg("max_new_tokens") = -1, py::arg("batch_size") = 512, py::arg("ctx_size") = 512, py::arg("seed") = -1,
            py::arg("threads") = 8, py::arg("repeat_penalty") = 1.1f, py::arg("post_process") = "topk")
-      .def("generate", &Model::generate, "Generate tokens with prompt", py::arg("prompt"),
-           py::arg("sentence_mode") = true)
+      .def("generate", &Model::generate, "Generate token with input ids", py::arg("input_ids"))
+      .def("generate_tokens", &Model::generate_tokens, "Generate tokens with input ids", py::arg("input_ids"))
       .def_static("quant_model", &Model::quant_model, "Quantize model", py::arg("model_path"), py::arg("out_path"),
                   py::arg("weight_dtype") = "int4", py::arg("alg") = "sym", py::arg("group_size") = 32,
                   py::arg("scale_dtype") = "fp32", py::arg("compute_dtype") = "ggml", py::arg("use_ggml") = false)

intel_extension_for_transformers/llm/runtime/graph/scripts/python_api_example.py

@@ -15,11 +15,17 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from transformers import AutoTokenizer, TextStreamer
 from intel_extension_for_transformers.transformers import AutoModel, WeightOnlyQuantConfig
-model_name = "mosaicml/mpt-7b"
-woq_config = WeightOnlyQuantConfig(compute_dtype="int8")
-
-model = AutoModel.from_pretrained(model_name, quantization_config=woq_config)
 
+model_name = "THUDM/chatglm2-6b"  # or local path to model
+woq_config = WeightOnlyQuantConfig(compute_dtype="int8", weight_dtype="int4")
 prompt = "Once upon a time, a little girl"
-print(model.generate(prompt, max_new_tokens=30))
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+inputs = tokenizer(prompt, return_tensors="pt").input_ids
+streamer = TextStreamer(tokenizer)
+
+model = AutoModel.from_pretrained(model_name, quantization_config=woq_config, trust_remote_code=True)
+outputs = model.generate(inputs, streamer=streamer, max_new_tokens=300)
+

intel_extension_for_transformers/neural_chat/examples/talkingbot_pc/build_talkingbot_on_pc.ipynb

@@ -99,13 +99,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from transformers import AutoTokenizer, TextStreamer\n",
     "from intel_extension_for_transformers.llm.runtime.graph import Model\n",
-    "model = Model()\n",
-    "model.bin_file = r\"mpt_q4_0.bin\"\n",
-    "model.init_from_bin(\"mpt\", model.bin_file, max_new_tokens=32, seed=12)\n",
+    "\n",
     "prompt = text\n",
-    "output = model.generate(prompt)\n",
-    "print(output)"
+    "tokenizer = AutoTokenizer.from_pretrained(\"EleutherAI/gpt-j-6b\", trust_remote_code=True)\n",
+    "inputs = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
+    "streamer = TextStreamer(tokenizer)\n",
+    "\n",
+    "model = Model()\n",
+    "model.init_from_bin(\"gptj\", \"ne_gptj_q.bin\", max_new_tokens=320, seed=12)\n",
+    "\n",
+    "outputs = model.generate(inputs, streamer=streamer)\n"
    ]
   },
   {
@@ -121,12 +126,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "from transformers import AutoTokenizer, TextStreamer\n",
     "from intel_extension_for_transformers.transformers import AutoModel, WeightOnlyQuantConfig\n",
-    "model_name = r\"THUDM/ChatGLM2-6B\"\n",
+    "\n",
+    "model_name = \"EleutherAI/gpt-j-6b\"  # or local path to model\n",
     "woq_config = WeightOnlyQuantConfig(compute_dtype=\"int8\", weight_dtype=\"int4\")\n",
-    "model = AutoModel.from_pretrained(model_name, quantization_config=woq_config, use_llm_runtime=True, trust_remote_code=True)\n",
     "prompt = text\n",
-    "output = model.generate(prompt, max_new_tokens=32)"
+    "\n",
+    "tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)\n",
+    "inputs = tokenizer(prompt, return_tensors=\"pt\").input_ids\n",
+    "streamer = TextStreamer(tokenizer)\n",
+    "\n",
+    "model = AutoModel.from_pretrained(model_name, quantization_config=woq_config, trust_remote_code=True)\n",
+    "outputs = model.generate(inputs, streamer=streamer, max_new_tokens=300)\n"
    ]
   },
   {

intel_extension_for_transformers/transformers/modeling/modeling_auto.py

@@ -143,7 +143,8 @@ def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
                     alg=quantization_config.scheme,
                     group_size=quantization_config.group_size,
                     scale_dtype=quantization_config.scale_dtype,
-                    compute_dtype=quantization_config.compute_dtype
+                    compute_dtype=quantization_config.compute_dtype,
+                    use_ggml=quantization_config.use_ggml,
                 )
                 return model
             else:

intel_extension_for_transformers/transformers/utils/quantization_config.py

@@ -39,6 +39,7 @@ def __init__(
         group_size=32,
         scheme="sym",
         algorithm="RTN",
+        use_ggml=False,
         **kwargs,
     ):
         from intel_extension_for_transformers.llm.quantization.utils import convert_dtype_2_str
@@ -57,6 +58,7 @@ def __init__(
         self.calib_dataset = kwargs.pop("calib_dataset", "NeelNanda/pile-10k")
         self.calib_dataloader = kwargs.pop("calib_dataloader", None)
         self.calib_iters = kwargs.pop("calib_iters", 100)
+        self.use_ggml = use_ggml
 
         if compute_dtype is None:
             self.compute_dtype = "fp32"
@@ -116,8 +118,8 @@ def post_init_runtime(self):
 
         if self.compute_dtype is None:
             self.compute_dtype = "int8"
-        elif self.compute_dtype not in ['int8', 'fp32']:
-            raise ValueError("compute_dtype must be 'int8', 'fp32'.")
+        elif self.compute_dtype not in ['int8', 'bf16', 'fp32']:
+            raise ValueError("compute_dtype must be 'int8', 'bf16', 'fp32'.")
 
         if self.weight_dtype is None:
             self.weight_dtype = "int4"

setup.py

@@ -57,12 +57,11 @@ def check_env_flag(name: str, default: bool = False) -> bool:
 class CMakeExtension(Extension):
     """CMakeExtension class."""
 
-    def __init__(self, name, sourcedir="", lib_only=False, compile=True):
+    def __init__(self, name, sourcedir="", lib_only=False):
         """Init a CMakeExtension object."""
         Extension.__init__(self, name, sources=[])
         self.sourcedir = os.path.abspath(sourcedir)
         self.optional = lib_only  # we only deliver shared object but not as a python extension module
-        self.compile = compile
 
 
 class CMakeBuild(build_ext):
@@ -107,8 +106,6 @@ def get_source_files(self):
         return files
 
     def build_extension(self, ext: CMakeExtension) -> None:
-        if not ext.compile:
-            return
         # Must be in this form due to bug in .resolve() only fixed in Python 3.10+
         ext_fullpath = Path.cwd() / self.get_ext_fullpath(ext.name)
         extdir = ext_fullpath.parent.resolve()
@@ -248,18 +245,9 @@ def check_submodules():
         check_submodules()
         ext_modules.extend([
             CMakeExtension("intel_extension_for_transformers.neural_engine_py", "intel_extension_for_transformers/llm/runtime/deprecated/"),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.gptj_cpp", "intel_extension_for_transformers/llm/runtime/graph/"),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.falcon_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.gptneox_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.dolly_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.llama_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.mpt_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.starcoder_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.opt_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.bloom_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False),
-            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.chatglm2_cpp", "intel_extension_for_transformers/llm/runtime/graph/", compile=False)
+            CMakeExtension("intel_extension_for_transformers.llm.runtime.graph.Model", "intel_extension_for_transformers/llm/runtime/graph/"),
             ])
-        cmdclass={'build_ext': CMakeBuild}
+    cmdclass={'build_ext': CMakeBuild}
 
     setup(
         name="intel-extension-for-transformers",