add examples

README.md

@@ -27,3 +27,174 @@ Quantized with [OmniQuant](https://github.com/OpenGVLab/OmniQuant).
 
 ## Examples
 
+### Imports and Model Loading
+
+<details>
+  <summary>Expand</summary>
+  
+  ```python
+  import gc
+
+  import auto_gptq.nn_modules.qlinear.qlinear_cuda as qlinear_cuda
+  import auto_gptq.nn_modules.qlinear.qlinear_triton as qlinear_triton
+  import torch
+
+  from accelerate import (
+      init_empty_weights,
+      infer_auto_device_map,
+      load_checkpoint_in_model,
+  )
+  from tqdm import tqdm
+  from transformers import (
+      AutoConfig,
+      AutoModelForCausalLM,
+      AutoTokenizer,
+      pipeline,
+  )
+
+
+  def get_named_linears(model):
+      return {
+          name: module for name, module in model.named_modules()
+          if isinstance(module, torch.nn.Linear)
+      }
+
+
+  def set_module(model, name, module):
+      parent = model
+      levels = name.split('.')
+
+      for i in range(len(levels) - 1):
+          cur_name = levels[i]
+
+          if cur_name.isdigit():
+              parent = parent[int(cur_name)]
+          else:
+              parent = getattr(parent, cur_name)
+
+      setattr(parent, levels[-1], module)
+
+
+  def load_model(model_path):
+      # Based on: https://github.com/OpenGVLab/OmniQuant/blob/main/runing_quantized_mixtral_7bx8.ipynb
+
+      config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
+
+      if not hasattr(config, 'quantization_config'):
+          raise AttributeError(
+              f'No quantization info found in model config "{model_path}"'
+              f' (`quantization_config` section is missing).'
+          )
+
+      wbits = config.quantization_config['bits']
+      group_size = config.quantization_config['group_size']
+
+      # We are going to init an ordinary model and then manually replace all Linears with QuantLinears
+      del config.quantization_config
+
+      with init_empty_weights():
+          model = AutoModelForCausalLM.from_config(config=config, torch_dtype=torch.float16, trust_remote_code=True)
+
+      layers = model.model.layers
+
+      for i in tqdm(range(len(layers))):
+          layer = layers[i]
+          named_linears = get_named_linears(layer)
+
+          for name, module in named_linears.items():
+              params = (
+                  wbits, group_size,
+                  module.in_features, module.out_features,
+                  module.bias is not None
+              )
+
+              if wbits in [2, 4]:
+                  q_linear = qlinear_triton.QuantLinear(*params)
+              elif wbits == 3:
+                  q_linear = qlinear_cuda.QuantLinear(*params)
+              else:
+                  raise NotImplementedError("Only 2, 3 and 4 bits are supported.")
+
+              q_linear.to(next(layer.parameters()).device)
+              set_module(layer, name, q_linear)
+
+      torch.cuda.empty_cache()
+      gc.collect()
+
+      model.tie_weights()
+      device_map = infer_auto_device_map(model)
+
+      print("Loading pre-computed quantized weights...")
+
+      load_checkpoint_in_model(
+          model, checkpoint=model_path,
+          device_map=device_map, offload_state_dict=True,
+      )
+
+      print("Model loaded successfully!")
+
+      return model
+  ```
+</details>
+
+
+### Inference
+
+```python
+model_path = "compressa-ai/Saiga-Llama-3-8B-OmniQuant"
+
+model = load_model(model_path)
+model.cuda()
+tokenizer = AutoTokenizer.from_pretrained(
+    model_path, use_fast=False, trust_remote_code=True
+)
+
+system_message = "Ты — дружелюбный чат-бот, который всегда отвечает как пират."
+user_message = "Куда мы направляемся, капитан?"
+messages = [
+    {"role": "system", "content": system_message},
+    {"role": "user", "content": user_message},
+]
+prompt = tokenizer.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+
+inputs = tokenizer(prompt, return_tensors="pt")
+inputs = {k: v.cuda() for k, v in inputs.items()}
+
+outputs = model.generate(
+    **inputs, max_new_tokens=512,
+    do_sample=True, temperature=0.7, top_p=0.95,
+)
+
+response = tokenizer.decode(outputs[0])  # , skip_special_tokens=True)
+continuation = response.removeprefix(prompt).removesuffix(tokenizer.eos_token)
+
+print(f'Prompt:\n{prompt}')
+print(f'Continuation:\n{continuation}\n')
+```
+
+
+### Inference Using Pipeline
+
+```python
+pipe = pipeline(
+    "text-generation",
+    model=model, tokenizer=tokenizer,
+    max_new_tokens=512, do_sample=True,
+    temperature=0.7, top_p=0.95,
+    device=0,
+)
+
+prompt = pipe.tokenizer.apply_chat_template(
+    messages, tokenize=False, add_generation_prompt=True
+)
+
+outputs = pipe(prompt)
+
+response = outputs[0]["generated_text"]
+continuation = response.removeprefix(prompt)
+
+print(f'Prompt:\n{prompt}')
+print(f'Continuation:\n{continuation}\n')
+```