From 099531c77b998bcfd14d641ebdb37fbc25fa9c5f Mon Sep 17 00:00:00 2001
From: pks <pks@pks.rocks>
Date: Sat, 6 Dec 2025 22:31:04 +0100
Subject: WIP

---
 inference.py | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/inference.py b/inference.py
index c522dd4..9f17c5f 100755
--- a/inference.py
+++ b/inference.py
@@ -90,12 +90,11 @@ def make_inputs_with_prefix(processor,
 def generate_and_parse(model,
                        processor,
                        messages,
-                       mode,
                        args,
                        example_id):
     sys.stderr.write(f"Processing {example_id=}\n")
 
-    if mode == "with_prefix":
+    if args.mode == "with_prefix":
         inputs, prefix = make_inputs_with_prefix(processor, messages, model.device)
     else:
         inputs = make_inputs(processor, messages, model.device)
@@ -117,8 +116,8 @@ def generate_and_parse(model,
 
     output_tokens = generation[0][input_len:]
     output_text = clean_str(processor.decode(output_tokens, skip_special_tokens=True))
-    if mode == "with_prefix":
-        output_text = (prefix + output_text).removeprefix("<bos><start_of_turn>userYou are a professional English-German translator and also a renowned photography critic.<start_of_image>Write a detailed caption for this image in a single sentence. Translate the caption into German. The output needs to be JSON, the keys being 'English' and 'German' for the respective captions. Only output the JSON, nothing else.<end_of_turn><start_of_turn>model")
+    if args.mode == "with_prefix":
+        output_text = prefix[prefix.index("{"):] + output_text
 
     try:
         return json.loads(output_text)
-- 
cgit v1.2.3