Upload code/inference.py with huggingface_hub

code/inference.py

@@ -45,17 +45,25 @@ def predict_fn(data, model_and_tokenizer):
     temperature = data.pop("temperature", 0.2)
     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 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)