Spaces:
Sleeping
Sleeping
| """ | |
| A model worker executes the model. | |
| """ | |
| import argparse | |
| import asyncio | |
| import json | |
| import time | |
| import threading | |
| import uuid | |
| import requests | |
| import torch | |
| from functools import partial | |
| from mplug_owl2.constants import WORKER_HEART_BEAT_INTERVAL | |
| from mplug_owl2.utils import (build_logger, server_error_msg, | |
| pretty_print_semaphore) | |
| from mplug_owl2.model.builder import load_pretrained_model | |
| from mplug_owl2.mm_utils import process_images, load_image_from_base64, tokenizer_image_token, KeywordsStoppingCriteria | |
| from mplug_owl2.constants import IMAGE_TOKEN_INDEX, DEFAULT_IMAGE_TOKEN | |
| from transformers import TextIteratorStreamer | |
| from threading import Thread | |
| GB = 1 << 30 | |
| worker_id = str(uuid.uuid4())[:6] | |
| logger = build_logger("model_worker", f"model_worker_{worker_id}.log") | |
| class ModelWorker: | |
| def __init__(self, model_path, model_base, model_name, load_8bit, load_4bit, device): | |
| self.worker_id = worker_id | |
| if model_path.endswith("/"): | |
| model_path = model_path[:-1] | |
| if model_name is None: | |
| model_paths = model_path.split("/") | |
| if model_paths[-1].startswith('checkpoint-'): | |
| self.model_name = model_paths[-2] + "_" + model_paths[-1] | |
| else: | |
| self.model_name = model_paths[-1] | |
| else: | |
| self.model_name = model_name | |
| self.device = device | |
| logger.info(f"Loading the model {self.model_name} on worker {worker_id} ...") | |
| self.tokenizer, self.model, self.image_processor, self.context_len = load_pretrained_model( | |
| model_path, model_base, self.model_name, load_8bit, load_4bit, device=self.device) | |
| self.is_multimodal = True | |
| def predict_stream(self, params): | |
| tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor | |
| prompt = params["prompt"] + "The quality of the image is" | |
| ori_prompt = prompt | |
| images = params.get("images", None) | |
| num_image_tokens = 0 | |
| if images is not None and len(images) > 0 and self.is_multimodal: | |
| if len(images) > 0: | |
| if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN): | |
| raise ValueError("Number of images does not match number of <|image|> tokens in prompt") | |
| images = [load_image_from_base64(image) for image in images] | |
| images = process_images(images, image_processor, model.config) | |
| if type(images) is list: | |
| images = [image.to(self.model.device, dtype=torch.float16) for image in images] | |
| else: | |
| images = images.to(self.model.device, dtype=torch.float16) | |
| replace_token = DEFAULT_IMAGE_TOKEN | |
| prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) | |
| num_image_tokens = prompt.count(replace_token) * (model.get_model().visual_abstractor.config.num_learnable_queries + 1) | |
| else: | |
| images = None | |
| image_args = {"images": images} | |
| else: | |
| images = None | |
| image_args = {} | |
| input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device) | |
| logits = model.forward( | |
| input_ids=input_ids, | |
| use_cache=True, | |
| **image_args).logits[0,-1] | |
| print(logits.shape) | |
| softmax_logits = torch.softmax(logits[[1781,6588,6460]], 0) | |
| print(tokenizer(["good", "average", "poor"])) | |
| fake_streamer = [] | |
| for id_, word in enumerate(["good", "average", "poor"]): | |
| stream_ = f"Probability of {word} quality: {softmax_logits[id_].item():.4f};\n" | |
| fake_streamer.append(stream_) | |
| quality_score = 0.5 * softmax_logits[1] + softmax_logits[0] | |
| stream_ = f"Quality score: {quality_score:.4f} (range [0,1])." | |
| fake_streamer.append(stream_) | |
| generated_text = ori_prompt.replace("The quality of the image is", "") | |
| for new_text in fake_streamer: | |
| generated_text += new_text | |
| yield json.dumps({"text": generated_text, "error_code": 0}).encode() | |
| def generate_stream(self, params): | |
| tokenizer, model, image_processor = self.tokenizer, self.model, self.image_processor | |
| prompt = params["prompt"] | |
| ori_prompt = prompt | |
| images = params.get("images", None) | |
| num_image_tokens = 0 | |
| if images is not None and len(images) > 0 and self.is_multimodal: | |
| if len(images) > 0: | |
| if len(images) != prompt.count(DEFAULT_IMAGE_TOKEN): | |
| raise ValueError("Number of images does not match number of <|image|> tokens in prompt") | |
| images = [load_image_from_base64(image) for image in images] | |
| images = process_images(images, image_processor, model.config) | |
| if type(images) is list: | |
| images = [image.to(self.model.device, dtype=torch.float16) for image in images] | |
| else: | |
| images = images.to(self.model.device, dtype=torch.float16) | |
| replace_token = DEFAULT_IMAGE_TOKEN | |
| prompt = prompt.replace(DEFAULT_IMAGE_TOKEN, replace_token) | |
| num_image_tokens = prompt.count(replace_token) * (model.get_model().visual_abstractor.config.num_learnable_queries + 1) | |
| else: | |
| images = None | |
| image_args = {"images": images} | |
| else: | |
| images = None | |
| image_args = {} | |
| temperature = float(params.get("temperature", 1.0)) | |
| top_p = float(params.get("top_p", 1.0)) | |
| max_context_length = getattr(model.config, 'max_position_embeddings', 4096) | |
| max_new_tokens = min(int(params.get("max_new_tokens", 256)), 1024) | |
| stop_str = params.get("stop", None) | |
| do_sample = True if temperature > 0.001 else False | |
| input_ids = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(self.device) | |
| keywords = [stop_str] | |
| stopping_criteria = KeywordsStoppingCriteria(keywords, tokenizer, input_ids) | |
| streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=15) | |
| max_new_tokens = min(max_new_tokens, max_context_length - input_ids.shape[-1] - num_image_tokens) | |
| if max_new_tokens < 1: | |
| yield json.dumps({"text": ori_prompt + "Exceeds max token length. Please start a new conversation, thanks.", "error_code": 0}).encode() + b"\0" | |
| return | |
| thread = Thread(target=model.generate, kwargs=dict( | |
| inputs=input_ids, | |
| do_sample=do_sample, | |
| temperature=temperature, | |
| top_p=top_p, | |
| max_new_tokens=max_new_tokens, | |
| streamer=streamer, | |
| stopping_criteria=[stopping_criteria], | |
| use_cache=True, | |
| **image_args | |
| )) | |
| thread.start() | |
| generated_text = ori_prompt | |
| for new_text in streamer: | |
| generated_text += new_text | |
| if generated_text.endswith(stop_str): | |
| generated_text = generated_text[:-len(stop_str)] | |
| yield json.dumps({"text": generated_text, "error_code": 0}).encode() | |
| def predict_stream_gate(self, params): | |
| try: | |
| for x in self.predict_stream(params): | |
| yield x | |
| except ValueError as e: | |
| print("Caught ValueError:", e) | |
| ret = { | |
| "text": server_error_msg, | |
| "error_code": 1, | |
| } | |
| yield json.dumps(ret).encode() | |
| except torch.cuda.CudaError as e: | |
| print("Caught torch.cuda.CudaError:", e) | |
| ret = { | |
| "text": server_error_msg, | |
| "error_code": 1, | |
| } | |
| yield json.dumps(ret).encode() | |
| except Exception as e: | |
| print("Caught Unknown Error", e) | |
| ret = { | |
| "text": server_error_msg, | |
| "error_code": 1, | |
| } | |
| yield json.dumps(ret).encode() | |
| def generate_stream_gate(self, params): | |
| try: | |
| for x in self.generate_stream(params): | |
| yield x | |
| except ValueError as e: | |
| print("Caught ValueError:", e) | |
| ret = { | |
| "text": server_error_msg, | |
| "error_code": 1, | |
| } | |
| yield json.dumps(ret).encode() | |
| except torch.cuda.CudaError as e: | |
| print("Caught torch.cuda.CudaError:", e) | |
| ret = { | |
| "text": server_error_msg, | |
| "error_code": 1, | |
| } | |
| yield json.dumps(ret).encode() | |
| except Exception as e: | |
| print("Caught Unknown Error", e) | |
| ret = { | |
| "text": server_error_msg, | |
| "error_code": 1, | |
| } | |
| yield json.dumps(ret).encode() |