Update __main__.py

__main__.py

@@ -1,4 +1,5 @@
 import argparse
+import sys
 import torch
 import torch.nn as nn
 from PIL import Image
@@ -12,26 +13,14 @@ from transformers import (
 from transformers import TextStreamer
 
 
-def tokenizer_image_token(
-    prompt, tokenizer, image_token_index=-200, return_tensors=None
-):
-    prompt_chunks = [tokenizer(chunk).input_ids for chunk in prompt.split("<image>")]
-
-    def insert_separator(X, sep):
-        return [ele for sublist in zip(X, [sep] * len(X)) for ele in sublist][:-1]
+def tokenizer_image_token(prompt, tokenizer, image_token_index=-200):
+    prompt_chunks = prompt.split("<image>")
+    tokenized_chunks = [tokenizer(chunk).input_ids for chunk in prompt_chunks]
+    input_ids = tokenized_chunks[0]
 
-    input_ids = []
-    offset = 0
-    if (
-        len(prompt_chunks) > 0
-        and len(prompt_chunks[0]) > 0
-        and prompt_chunks[0][0] == tokenizer.bos_token_id
-    ):
-        offset = 1
-        input_ids.append(prompt_chunks[0][0])
-
-    for x in insert_separator(prompt_chunks, [image_token_index] * (offset + 1)):
-        input_ids.extend(x[offset:])
+    for chunk in tokenized_chunks[1:]:
+        input_ids.append(image_token_index)
+        input_ids.extend(chunk[1:])  # Exclude BOS token on nonzero index
 
     return torch.tensor(input_ids, dtype=torch.long)
 
@@ -126,18 +115,20 @@ def answer_question(
     tokenizer.eos_token = "<|eot_id|>"
 
     try:
-        q = input("user: ")
+        q = input("\nuser: ")
     except EOFError:
         q = ""
     if not q:
         print("no input detected. exiting.")
+        sys.exit()
+
 
     question = "<image>" + q
 
-    prompt = f"<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    prompt = f"<|start_header_id|>user<|end_header_id|>\n\n{question}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
 
     input_ids = (
-        tokenizer_image_token(prompt, tokenizer, -200, return_tensors="pt")
+        tokenizer_image_token(prompt, tokenizer)
         .unsqueeze(0)
         .to(model.device)
     )
@@ -183,13 +174,14 @@ def answer_question(
         }
 
         while True:
+            print('assistant: ')
             generated_ids = model.generate(
                 inputs_embeds=new_embeds, attention_mask=attn_mask, **model_kwargs
             )[0]
 
             generated_text = tokenizer.decode(generated_ids, skip_special_tokens=False)
             try:
-                q = input("user: ")
+                q = input("\nuser: ")
             except EOFError:
                 q = ""
             if not q: