Browser tweaks, fix chat

Author: Varun Puri

api/api.py

@@ -235,78 +235,3 @@ def _callback(self, x, *, buffer, done_generating):
             done_generating = True
         if self.is_llama3_model and x.item() == self.tokenizer.special_tokens["<|eot_id|>"]:
             done_generating = True
-
-def main(args):
-    builder_args = BuilderArgs.from_args(args)
-    speculative_builder_args = BuilderArgs.from_speculative_args(args)
-    tokenizer_args = TokenizerArgs.from_args(args)
-    generator_args = GeneratorArgs.from_args(args)
-    generator_args.chat_mode = False
-
-    @st.cache_resource
-    def initialize_generator() -> OpenAIAPIGenerator:
-        return OpenAIAPIGenerator(
-            builder_args,
-            speculative_builder_args,
-            tokenizer_args,
-            generator_args,
-            args.profile,
-            args.quantize,
-            args.draft_quantize,
-        )
-
-    gen = initialize_generator()
-
-    tokens_generated = 0
-    st.title("Simple chat")
-
-    # Initialize chat history
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-
-    # Accept user input
-    if prompt := st.chat_input("What is up?"):
-        # Add user message to chat history
-        st.session_state.messages.append({"role": "user", "content": prompt})
-        # Display user message in chat message container
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
-        # Display assistant response in chat message container
-        with st.chat_message("assistant"):
-            
-            req = CompletionRequest(
-                model = gen.builder_args.checkpoint_path,
-                prompt = prompt,
-                temperature = generator_args.temperature,
-                messages = [],
-            )
-
-            def unwrap(completion_generator):
-                for chunk_response in completion_generator:
-                    content = chunk_response.choices[0].delta.content
-                    if not gen.is_llama3_model or content not in set(gen.tokenizer.special_tokens.keys()):
-                        yield content
-                yield "."
-                
-            response = st.write_stream(unwrap(gen.completion(req)))
-
-        # Add assistant response to chat history
-        st.session_state.messages.append({"role": "assistant", "content": response})
-
-
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="torchchat generate CLI")
-    verb = "generate"
-    add_arguments_for_verb(parser, verb)
-    args = parser.parse_args()
-    check_args(args, verb)
-    args = arg_init(args)
-    main(args)

api/browser.py

@@ -38,7 +38,7 @@ def initialize_generator() -> OpenAIAPIGenerator:
     gen = initialize_generator()
 
     tokens_generated = 0
-    st.title("Simple chat")
+    st.title("TorchChat")
 
     # Initialize chat history
     if "messages" not in st.session_state:
@@ -58,7 +58,7 @@ def initialize_generator() -> OpenAIAPIGenerator:
             st.markdown(prompt)
 
         # Display assistant response in chat message container
-        with st.chat_message("assistant"):
+        with st.chat_message("assistant"), st.status("Generating... ", expanded=True) as status:
 
             req = CompletionRequest(
                 model = gen.builder_args.checkpoint_path,
@@ -68,11 +68,16 @@ def initialize_generator() -> OpenAIAPIGenerator:
             )
 
             def unwrap(completion_generator):
+                start = time.time()
+                tokcount = 0
                 for chunk_response in completion_generator:
                     content = chunk_response.choices[0].delta.content
                     if not gen.is_llama3_model or content not in set(gen.tokenizer.special_tokens.keys()):
                         yield content
-                yield "."
+                    if content == gen.tokenizer.eos_id():
+                        yield "."
+                    tokcount+=1
+                status.update(label="Done, averaged {:.2f} tokens/second".format(tokcount/(time.time()-start)), state="complete")
 
             response = st.write_stream(unwrap(gen.completion(req)))
 

generate.py

@@ -823,7 +823,8 @@ def main(args):
         args.quantize,
         args.draft_quantize,
     )
-    gen.chat(generator_args)
+    for _ in gen.chat(generator_args):
+        pass
 
 
 if __name__ == "__main__":