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,31 @@ 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")
+
+    if conv_mode == "raw":
+        # use raw_prompt as prompt
+        prompt = raw_prompt
+        stop_str = "###"
+    else:
+        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)