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")
 
 
@@ -26,15 +27,18 @@ def reset_chat():
     st.session_state["messages"] = start_state
     st.session_state["conversation_images"] = []
 
+
 if "messages" not in st.session_state:
     st.session_state.messages = start_state
 if "conversation_images" not in st.session_state:
     st.session_state.conversation_images = []
 
+
 def _upload_image_prompts(file_uploads):
     for file in file_uploads:
         st.session_state.conversation_images.append(file)
 
+
 with st.sidebar:
     if st.button("Reset Chat", type="primary"):
         reset_chat()
@@ -105,7 +109,6 @@ def _upload_image_prompts(file_uploads):
         for img in st.session_state.conversation_images:
             st.image(img)
     st.session_state.conversation_images = []
-        
 
     with st.chat_message("assistant"), st.status(
         "Generating... ", expanded=True