Merge pull request xusenlinzy#313 from white2018/cogvlm2

support cogvlm2 model

Author: xusenlinzy

api/adapter/patcher.py

@@ -164,6 +164,9 @@ def patch_config(
 
 
 def patch_model(model: "PreTrainedModel") -> None:
+ if 'CogVLMForCausalLM' in model.config.architectures:
+ model.config.model_type = "cogvlm2"
+ return
  if model.config.model_type == "internvl_chat":
  return
  if model.config.model_type == "minicpmv":

api/engine/hf.py

@@ -37,6 +37,7 @@
 from api.templates.glm import generate_stream_chatglm, generate_stream_chatglm_v3
 from api.templates.minicpm import generate_stream_minicpm_v
 from api.templates.minimonkey import generate_stream_minimonkey
+from api.templates.cogvlm2 import generate_stream_cogvlm2
 from api.templates.stream import generate_stream
 from api.templates.utils import get_context_length
 from api.utils import create_error_response
@@ -81,6 +82,8 @@ def __init__(
  self.generate_stream_func = generate_stream_minicpm_v
  elif self.model.config.model_type == "internvl_chat":
  self.generate_stream_func = generate_stream_minimonkey
+ elif self.model.config.model_type == "cogvlm2":
+ self.generate_stream_func = generate_stream_cogvlm2
 
  logger.info(f"Using {self.model_name} Model for Chat!")
  logger.info(f"Using {self.template} for Chat!")
@@ -101,6 +104,8 @@ def _generate(self, params: Dict[str, Any]) -> Iterator[dict]:
  else:
  if self.model.config.model_type == "minicpmv":
  inputs = prompt_or_messages
+ elif self.model.config.model_type == "cogvlm2":
+ inputs = prompt_or_messages
  elif self.model.config.model_type == "internvl_chat":
  inputs = prompt_or_messages
  else:

api/templates/cogvlm2.py

@@ -0,0 +1,184 @@
+from __future__ import annotations
+
+import gc
+import time
+import uuid
+from typing import (
+ Any,
+ Dict,
+ List,
+ Iterator,
+ TYPE_CHECKING,
+)
+
+import torch
+
+from api.protocol import ChatCompletionMessageParam
+
+if TYPE_CHECKING:
+ from transformers import PreTrainedTokenizer, PreTrainedModel
+
+
+import queue
+from threading import Thread
+import torchvision.transforms as T
+import transformers
+from torchvision.transforms.functional import InterpolationMode
+from transformers import BitsAndBytesConfig, TextIteratorStreamer
+
+transformers.logging.set_verbosity_error()
+
+# THUDM/cogvlm2-llama3-chat-19B
+# THUDM/cogvlm2-llama3-chinese-chat-19B
+
+@torch.inference_mode()
+def generate_stream_cogvlm2(
+ model: "PreTrainedModel",
+ tokenizer: "PreTrainedTokenizer",
+ params: Dict[str, Any],
+) -> Iterator:
+ """
+ Generates text in a streaming manner using the ChatGLM model.
+
+ Args:
+ model: The pre-trained model.
+ tokenizer: The tokenizer used for tokenizing the input.
+ params: A dictionary containing the input parameters.
+
+ Yields:
+ A dictionary representing each generated text completion.
+
+ """
+ inputs = params["inputs"]
+ model_name = params.get("model", "llm")
+
+ query, history, images, system_message = prompt_history_images_system_from_messages(inputs, img_tok='')
+
+ input_by_model = model.build_conversation_input_ids(tokenizer, query=query, history=history, images=images, template_version='chat')
+
+ inputs = {
+ 'input_ids': input_by_model['input_ids'].unsqueeze(0).to(model.device),
+ 'token_type_ids': input_by_model['token_type_ids'].unsqueeze(0).to(model.device),
+ 'attention_mask': input_by_model['attention_mask'].unsqueeze(0).to(model.device),
+ 'images': [[input_by_model['images'][0].to(model.device).to(model.dtype)]] if images else None,
+ }
+
+ new_params = dict(temperature = float(params.get("temperature", 1.0)),
+ max_new_tokens = int(params.get("max_tokens", 256)),
+ repetition_penalty = float(params.get("repetition_penalty", 1.0)),
+ top_p = float(params.get("top_p", 1.0)),
+ top_k = int(params.get("top_k", 50)))
+
+ generation_kwargs = dict(
+ **inputs,
+ **new_params,
+ )
+
+ input_echo_len = 0
+ generated_text, previous_text = "", ""
+ completion_id: str = f"cmpl-{str(uuid.uuid4())}"
+ created: int = int(time.time())
+ for i, new_text in enumerate(threaded_streaming_generator(generate=model.generate, tokenizer=tokenizer, generation_kwargs=generation_kwargs)):
+ end = new_text.find(tokenizer.eos_token)
+ if end != -1:
+ new_text = new_text[:end]
+
+ generated_text += new_text
+ delta_text = generated_text[len(previous_text):]
+ previous_text = generated_text
+ yield {
+ "id": completion_id,
+ "object": "text_completion",
+ "created": created,
+ "model": model_name,
+ "delta": delta_text,
+ "text": generated_text,
+ "logprobs": None,
+ "finish_reason": None,
+ "usage": {
+ "prompt_tokens": input_echo_len,
+ "completion_tokens": i,
+ "total_tokens": input_echo_len + i,
+ },
+ }
+
+ if end != -1:
+ break
+
+ gc.collect()
+ torch.cuda.empty_cache()
+
+def prompt_history_images_system_from_messages(messages: list[ChatCompletionMessageParam], img_tok = "<image>\n"):
+ history = []
+ images = []
+ prompt = ''
+ system_prompt = None
+
+ for m in messages:
+ if m['role'] == 'user':
+ p = ''
+ for c in m['content']:
+ if c['type'] == 'image_url':
+ image = url_to_image(c['image_url']['url'])
+ images.extend([image])
+ p = img_tok + p
+ if c['type'] == 'text':
+ p += c['text']
+
+ prompt += p
+ elif m['role'] == 'assistant':
+ for c in m['content']:
+ if c['type'] == 'text':
+ history.extend([(prompt, c['text'])])
+ prompt = ''
+ elif m['role'] == 'system':
+ for c in m['content']:
+ if c['type'] == 'text':
+ system_prompt = c['text']
+
+ return prompt, history, images, system_prompt
+
+
+def url_to_image(image_url: str):
+ from PIL import Image
+ from io import BytesIO
+
+ if image_url.startswith("data:"):
+ import base64
+
+ image_bytes = base64.b64decode(image_url.split(",")[1])
+ else:
+ import urllib.request
+
+ with urllib.request.urlopen(image_url) as f:
+ image_bytes = f.read()
+
+ return Image.open(BytesIO(image_bytes)).convert("RGB")
+
+
+def threaded_streaming_generator(generate, tokenizer, generation_kwargs):
+ streamer = TextIteratorStreamer(tokenizer, skip_special_tokens=True, skip_prompt=True, timeout=60)
+
+ generation_kwargs['streamer'] = streamer
+
+ exq = queue.Queue()
+
+ def wrapper():
+ try:
+ with torch.no_grad():
+ generate(**generation_kwargs)
+
+ except Exception as e:
+ #logger.exception(e)
+ exq.put(e)
+ streamer.end()
+
+ t = Thread(target=wrapper, daemon=True)
+ t.start()
+
+ for text in streamer:
+ if text:
+ yield text
+
+ if not exq.empty():
+ raise exq.get_nowait()