remove duplicate code in generate + fix generate.py prompt for llama2 (#590)

* remove duplicate code in generate

* fix indent

* add fixes

Author: metascroy

generate.py

@@ -463,6 +463,24 @@ def get_device_info(name: str) -> str:
         return torch.cuda.get_device_name(0)
     return ""
 
+def _callback(x, buffer, period_id, done_generating, tokenizer, is_llama3_model):
+    if done_generating:
+        return
+    buffer.append(
+        tokenizer.decode([period_id] + x.tolist())[1:]
+    )  # I think this results in the first output token being dropped from the display which is wrong.
+    if x.item() == tokenizer.eos_id():
+        done_generating = True
+    if (
+        is_llama3_model
+        and x.item() == tokenizer.special_tokens["<|eot_id|>"]
+    ):
+        done_generating = True
+        buffer = buffer[:-1]  # drop the eot_id from the output buffer
+    if len(buffer) == 4 or done_generating:
+        print("".join(buffer), end="", flush=True)
+        buffer.clear()
+    # print(, end='', flush=True)
 
 def _main(
     builder_args: BuilderArgs,
@@ -612,7 +630,7 @@ def _main(
                 break
             if not is_llama3_model:
                 if system_prompt:
-                    prompt = f"{B_INST} {B_SYS}\n{system_prompt.strip()}\n{E_SYS}\n\n{prompt.strip} {E_INST}"
+                    prompt = f"{B_INST} {B_SYS}\n{system_prompt.strip()}\n{E_SYS}\n\n{prompt.strip()} {E_INST}"
                     system_prompt = (
                         None  # can only provide system prompt on first interaction
                     )
@@ -659,53 +677,17 @@ def _main(
             period_id = tokenizer.encode(".")[0]
             done_generating = False
 
-            def callback(
-                x, buffer=buffer, period_id=period_id, done_generating=done_generating
-            ):
-                if done_generating:
-                    return
-                buffer.append(
-                    tokenizer.decode([period_id] + x.tolist())[1:]
-                )  # I think this results in the first output token being dropped from the display which is wrong.
-                if x.item() == tokenizer.eos_id():
-                    done_generating = True
-                if (
-                    is_llama3_model
-                    and x.item() == tokenizer.special_tokens["<|eot_id|>"]
-                ):
-                    done_generating = True
-                    buffer = buffer[:-1]  # drop the eot_id from the output buffer
-                if len(buffer) == 4 or done_generating:
-                    print("".join(buffer), end="", flush=True)
-                    buffer.clear()
-                # print(, end='', flush=True)
+            def callback(x):
+                return _callback(x, buffer=buffer, period_id=period_id, done_generating=done_generating, tokenizer=tokenizer, is_llama3_model=is_llama3_model)
 
         else:
             assert not generator_args.chat_mode
             buffer = [generator_args.prompt]
             period_id = tokenizer.encode(".")[0]
             done_generating = False
 
-            def callback(
-                x, buffer=buffer, period_id=period_id, done_generating=done_generating
-            ):
-                if done_generating:
-                    return
-                buffer.append(
-                    tokenizer.decode([period_id] + x.tolist())[1:]
-                )  # I think this results in the first output token being dropped from the display which is wrong.
-                if x.item() == tokenizer.eos_id():
-                    done_generating = True
-                if (
-                    is_llama3_model
-                    and x.item() == tokenizer.special_tokens["<|eot_id|>"]
-                ):
-                    done_generating = True
-                    buffer = buffer[:-1]  # drop the eot_id from the output buffer
-                if len(buffer) == 4 or done_generating:
-                    print("".join(buffer), end="", flush=True)
-                    buffer.clear()
-                # print(, end='', flush=True)
+            def callback(x):
+                return _callback(x, buffer=buffer, period_id=period_id, done_generating=done_generating, tokenizer=tokenizer, is_llama3_model=is_llama3_model)
 
         t0 = time.perf_counter()
         import contextlib