Marlon Wiprud

fix: use diff path

19f8931 almost 3 years ago

4.69 kB

from typing import Dict, List, Any
from transformers import pipeline
from PIL import Image
import requests
from transformers import AutoModelForCausalLM, LlamaTokenizer
import torch
from accelerate import (
    init_empty_weights,
    infer_auto_device_map,
    load_checkpoint_and_dispatch,
)


class EndpointHandler:
    def __init__(self, path=""):
        # Preload all the elements you are going to need at inference.

        # self.pipeline = pipeline(
        #     "text-generation", model="THUDM/cogvlm-chat-hf", trust_remote_code=True
        # )

        # self.model = AutoModelForCausalLM.from_pretrained(
        #     "THUDM/cogvlm-chat-hf", trust_remote_code=True
        # )

        self.tokenizer = LlamaTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5")

        # self.model = (
        #     AutoModelForCausalLM.from_pretrained(
        #         "THUDM/cogvlm-chat-hf",
        #         torch_dtype=torch.bfloat16,
        #         low_cpu_mem_usage=True,
        #         trust_remote_code=True,
        #     )
        #     .to("cuda")
        #     .eval()
        # )

        # DISTRIBUTED GPUS
        with init_empty_weights():
            self.model = AutoModelForCausalLM.from_pretrained(
                "THUDM/cogvlm-chat-hf",
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
                trust_remote_code=True,
            )

        device_map = infer_auto_device_map(
            self.model,
            max_memory={
                0: "16GiB",
                1: "16GiB",
                2: "16GiB",
                3: "16GiB",
                "cpu": "180GiB",
            },
            no_split_module_classes=["CogVLMDecoderLayer"],
        )
        self.model = load_checkpoint_and_dispatch(
            self.model,
            "~/.cache/huggingface/modules/transformers_modules/THUDM/cogvlm-chat-hf/8abca878c4257412c4c38eeafaed3fe27a036730",  # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
            # "/home/ec2-user/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/8abca878c4257412c4c38eeafaed3fe27a036730",  # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
            device_map=device_map,
            no_split_module_classes=["CogVLMDecoderLayer"],
        )
        self.model = self.model.eval()
        ## DISTRIBUTED GPUS

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
         data args:
              inputs (:obj: `str` | `PIL.Image` | `np.array`)
              kwargs
        Return:
              A :obj:`list` | `dict`: will be serialized and returned
        """

        query = data["query"]
        img_uri = data["img_uri"]

        image = Image.open(
            requests.get(
                img_uri,
                stream=True,
            ).raw
        ).convert("RGB")

        inputs = self.model.build_conversation_input_ids(
            self.tokenizer,
            query=query,
            history=[],
            images=[image],
            template_version="vqa",
        )  # vqa mode

        inputs = {
            "input_ids": inputs["input_ids"].unsqueeze(0).to("cuda"),
            "token_type_ids": inputs["token_type_ids"].unsqueeze(0).to("cuda"),
            "attention_mask": inputs["attention_mask"].unsqueeze(0).to("cuda"),
            "images": [[inputs["images"][0].to("cuda").to(torch.bfloat16)]],
        }

        gen_kwargs = {"max_length": 2048, "do_sample": False}

        with torch.no_grad():
            outputs = self.model.generate(**inputs, **gen_kwargs)
            outputs = outputs[:, inputs["input_ids"].shape[1] :]
            response = self.tokenizer.decode(outputs[0])
            return response


# query = "How many houses are there in this cartoon?"
# image = Image.open(
#     requests.get(
#         "https://github.com/THUDM/CogVLM/blob/main/examples/3.jpg?raw=true", stream=True
#     ).raw
# ).convert("RGB")
# inputs = model.build_conversation_input_ids(
#     tokenizer, query=query, history=[], images=[image], template_version="vqa"
# )  # vqa mode
# inputs = {
#     "input_ids": inputs["input_ids"].unsqueeze(0).to("cuda"),
#     "token_type_ids": inputs["token_type_ids"].unsqueeze(0).to("cuda"),
#     "attention_mask": inputs["attention_mask"].unsqueeze(0).to("cuda"),
#     "images": [[inputs["images"][0].to("cuda").to(torch.bfloat16)]],
# }
# gen_kwargs = {"max_length": 2048, "do_sample": False}

# with torch.no_grad():
#     outputs = model.generate(**inputs, **gen_kwargs)
#     outputs = outputs[:, inputs["input_ids"].shape[1] :]
#     print(tokenizer.decode(outputs[0]))

	from typing import Dict, List, Any
	from transformers import pipeline
	from PIL import Image
	import requests
	from transformers import AutoModelForCausalLM, LlamaTokenizer
	import torch
	from accelerate import (
	init_empty_weights,
	infer_auto_device_map,
	load_checkpoint_and_dispatch,
	)


	class EndpointHandler:
	def __init__(self, path=""):
	# Preload all the elements you are going to need at inference.

	# self.pipeline = pipeline(
	# "text-generation", model="THUDM/cogvlm-chat-hf", trust_remote_code=True
	# )

	# self.model = AutoModelForCausalLM.from_pretrained(
	# "THUDM/cogvlm-chat-hf", trust_remote_code=True
	# )

	self.tokenizer = LlamaTokenizer.from_pretrained("lmsys/vicuna-7b-v1.5")

	# self.model = (
	# AutoModelForCausalLM.from_pretrained(
	# "THUDM/cogvlm-chat-hf",
	# torch_dtype=torch.bfloat16,
	# low_cpu_mem_usage=True,
	# trust_remote_code=True,
	# )
	# .to("cuda")
	# .eval()
	# )

	# DISTRIBUTED GPUS
	with init_empty_weights():
	self.model = AutoModelForCausalLM.from_pretrained(
	"THUDM/cogvlm-chat-hf",
	torch_dtype=torch.bfloat16,
	low_cpu_mem_usage=True,
	trust_remote_code=True,
	)

	device_map = infer_auto_device_map(
	self.model,
	max_memory={
	0: "16GiB",
	1: "16GiB",
	2: "16GiB",
	3: "16GiB",
	"cpu": "180GiB",
	},
	no_split_module_classes=["CogVLMDecoderLayer"],
	)
	self.model = load_checkpoint_and_dispatch(
	self.model,
	"~/.cache/huggingface/modules/transformers_modules/THUDM/cogvlm-chat-hf/8abca878c4257412c4c38eeafaed3fe27a036730", # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
	# "/home/ec2-user/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/8abca878c4257412c4c38eeafaed3fe27a036730", # typical, '~/.cache/huggingface/hub/models--THUDM--cogvlm-chat-hf/snapshots/balabala'
	device_map=device_map,
	no_split_module_classes=["CogVLMDecoderLayer"],
	)
	self.model = self.model.eval()
	## DISTRIBUTED GPUS

	def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
	"""
	data args:
	inputs (:obj: `str` \| `PIL.Image` \| `np.array`)
	kwargs
	Return:
	A :obj:`list` \| `dict`: will be serialized and returned
	"""

	query = data["query"]
	img_uri = data["img_uri"]

	image = Image.open(
	requests.get(
	img_uri,
	stream=True,
	).raw
	).convert("RGB")

	inputs = self.model.build_conversation_input_ids(
	self.tokenizer,
	query=query,
	history=[],
	images=[image],
	template_version="vqa",
	) # vqa mode

	inputs = {
	"input_ids": inputs["input_ids"].unsqueeze(0).to("cuda"),
	"token_type_ids": inputs["token_type_ids"].unsqueeze(0).to("cuda"),
	"attention_mask": inputs["attention_mask"].unsqueeze(0).to("cuda"),
	"images": [[inputs["images"][0].to("cuda").to(torch.bfloat16)]],
	}

	gen_kwargs = {"max_length": 2048, "do_sample": False}

	with torch.no_grad():
	outputs = self.model.generate(inputs, gen_kwargs)
	outputs = outputs[:, inputs["input_ids"].shape[1] :]
	response = self.tokenizer.decode(outputs[0])
	return response


	# query = "How many houses are there in this cartoon?"
	# image = Image.open(
	# requests.get(
	# "https://github.com/THUDM/CogVLM/blob/main/examples/3.jpg?raw=true", stream=True
	# ).raw
	# ).convert("RGB")
	# inputs = model.build_conversation_input_ids(
	# tokenizer, query=query, history=[], images=[image], template_version="vqa"
	# ) # vqa mode
	# inputs = {
	# "input_ids": inputs["input_ids"].unsqueeze(0).to("cuda"),
	# "token_type_ids": inputs["token_type_ids"].unsqueeze(0).to("cuda"),
	# "attention_mask": inputs["attention_mask"].unsqueeze(0).to("cuda"),
	# "images": [[inputs["images"][0].to("cuda").to(torch.bfloat16)]],
	# }
	# gen_kwargs = {"max_length": 2048, "do_sample": False}

	# with torch.no_grad():
	# outputs = model.generate(inputs, gen_kwargs)
	# outputs = outputs[:, inputs["input_ids"].shape[1] :]
	# print(tokenizer.decode(outputs[0]))