Include empty assistant message for chat

Author: vmpuri

torchchat/generate.py

@@ -603,6 +603,7 @@ def generate(
         if len(prompt.shape) > 1:
             prompt = prompt.squeeze(0)
         prompt_length = prompt.size(0)
+        max_new_tokens = min(max_new_tokens, max_seq_length - start_pos - prompt_length)
         # set up caches only if first inference
         if start_pos == 0:
             model = model.to(device=device)
@@ -825,6 +826,12 @@ def _gen_model_input(
                         content=content,
                     )
                 )
+        messages.append(
+            Message(
+                role="assistant",
+                content="",
+            )
+        )
 
         transform = llama3_2_vision_transform(str(self.tokenizer_args.tokenizer_path))
 
@@ -849,7 +856,7 @@ def _gen_model_input(
                 seq_len = encoded.size(0)
                 batch = {}
 
-            total_response_length = max_seq_len + max_new_tokens
+            total_response_length = seq_len + max_new_tokens
             batch["causal_mask"] = torch.nn.functional.pad(
                 torch.tril(
                     torch.ones(

torchchat/usages/browser.py

@@ -10,6 +10,7 @@
 
 from openai import OpenAI
 
+st.set_page_config(page_title="torchchat", page_icon="🤖")
 st.title("torchchat")