Upload code/inference.py with huggingface_hub

code/inference.py

@@ -6,9 +6,16 @@ from transformers import AutoTokenizer
 
 from llava.model import LlavaLlamaForCausalLM
 from llava.utils import disable_torch_init
-from llava.constants import IMAGE_TOKEN_INDEX
 from llava.mm_utils import tokenizer_image_token, KeywordsStoppingCriteria
 
+from llava.conversation import conv_templates, SeparatorStyle
+from llava.constants import (
+    IMAGE_TOKEN_INDEX,
+    DEFAULT_IMAGE_TOKEN,
+    DEFAULT_IM_START_TOKEN,
+    DEFAULT_IM_END_TOKEN,
+)
+
 
 def model_fn(model_dir):
     kwargs = {"device_map": "auto"}
@@ -32,11 +39,23 @@ def predict_fn(data, model_and_tokenizer):
 
     # get prompt & parameters
     image_file = data.pop("image", data)
-    prompt = data.pop("question", data)
+    raw_prompt = data.pop("question", data)
 
     max_new_tokens = data.pop("max_new_tokens", 1024)
     temperature = data.pop("temperature", 0.2)
-    stop_str = data.pop("stop_str", "###")
+    conv_mode = data.pop("conv_mode", "llava_v1")
+
+    # conv_mode = "llava_v1"
+    conv = conv_templates[conv_mode].copy()
+    roles = conv.roles
+    inp = f"{roles[0]}: {raw_prompt}"
+    inp = (
+        DEFAULT_IM_START_TOKEN + DEFAULT_IMAGE_TOKEN + DEFAULT_IM_END_TOKEN + "\n" + inp
+    )
+    conv.append_message(conv.roles[0], inp)
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt()
+    stop_str = conv.sep if conv.sep_style != SeparatorStyle.TWO else conv.sep2
 
     if image_file.startswith("http") or image_file.startswith("https"):
         response = requests.get(image_file)