Update __main__.py

__main__.py

@@ -1,4 +1,6 @@
 import argparse
+import sys
+
 import numpy as np
 import torch
 import torch.nn as nn
@@ -120,7 +122,7 @@ def answer_question(
     except EOFError:
         inp = ""
     if not inp:
-        print("exit...")
+        sys.exit("exiting..")
 
     question = '<image>' + inp
 
@@ -163,6 +165,8 @@ def answer_question(
         }
 
         while True:
+            print('assistant: ')
+            
             generated_ids = model.generate(
                 inputs_embeds=new_embeds,
                 attention_mask=attn_mask,
@@ -176,7 +180,8 @@ def answer_question(
             except EOFError:
                 inp = ""
             if not inp:
-                print("exit...")
+                print("exiting...")
+                break
 
             new_text = generated_text + "<|start_header_id|>user<|end_header_id|>\n\n" + inp + "<|start_header_id|>assistant<|end_header_id|>\n\n"
             new_input_ids = tokenizer(new_text, return_tensors='pt').input_ids.to(device)
@@ -201,4 +206,4 @@ if __name__ == "__main__":
         vision_model,
         processor,
         projection_module,
-    )
+    )