Spaces:
Sleeping
Sleeping
File size: 2,323 Bytes
ba9f758 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from transformers import AutoModel, AutoTokenizer
import torch
from payload_model import PayloadModel
from internvl_utils import load_image
from pydantic import BaseModel, Field
from typing import Optional
class InternVL3(BaseModel):
model_name: str
model: Optional[AutoModel] = None
tokenizer: Optional[AutoTokenizer] = None
generation_config: dict = Field(default_factory=lambda: {"max_new_tokens": 1024, "do_sample": True})
model_config = {
"arbitrary_types_allowed": True,
"from_attributes": True
}
def __init__(self, model_name: str, **kwargs):
super().__init__(model_name=model_name, **kwargs)
self.model = AutoModel.from_pretrained(
self.model_name,
torch_dtype=torch.bfloat16,
load_in_8bit=False,
low_cpu_mem_usage=True,
use_flash_attn=True,
trust_remote_code=True,
device_map="auto",
).eval()
self.tokenizer = AutoTokenizer.from_pretrained(
self.model_name,
trust_remote_code=True,
use_fast=False,
)
def get_query_prompt(self, prompt_keyword: str):
if prompt_keyword == "person_running":
query_prompt = """
<image>\nCheck if person is running or not? If they are running
respond with "Yes" else respond with "No". Limit your response to either "Yes" or "No"
"""
else:
query_prompt = None
return query_prompt
def predict(self, payload: PayloadModel):
pixel_values = load_image(payload.image)
query_prompt = self.get_query_prompt(payload.prompt_keyword)
if query_prompt is None:
model_response = f"Invalid prompt keyword: {payload.prompt_keyword}"
else:
model_response = self.model.chat(
self.tokenizer,
pixel_values,
query_prompt,
generation_config=self.generation_config,
)
return model_response
def extract_model_response(self, model_response: str):
return "Yes" in model_response
async def __call__(self, payload: PayloadModel):
model_response = self.predict(payload)
extracted_response = self.extract_model_response(model_response)
return extracted_response
|