Replace browser UI with Basic Streamlit UI Implementation

Author: Varun Puri

browser/browser.py

@@ -1,37 +1,97 @@
-import subprocess
-import sys
+
+import streamlit as st
+
+from build.builder import (
+    _initialize_model,
+    _initialize_tokenizer,
+    BuilderArgs,
+    TokenizerArgs,
+)
+from build.model import Transformer
+from build.utils import device_sync, set_precision
+from cli import add_arguments_for_verb, arg_init, check_args, logger
+
+from generate import GeneratorArgs
+from api.api import *
+
 
 
 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("TorchChat")
+
+    # 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"), st.status("Generating... ", expanded=True) as status:
+            
+            req = CompletionRequest(
+                model = gen.builder_args.checkpoint_path,
+                prompt = prompt,
+                temperature = generator_args.temperature,
+                messages = [],
+            )
+
+            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
+                    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)))
+
+        # Add assistant response to chat history
+        st.session_state.messages.append({"role": "assistant", "content": response})
+
+
+
 
-    # Directory Containing the server file "chat_in_browser.py"
-    server_dir = "browser"
-
-    # Look for port from cmd args. Default to 5000 if not found.
-    port = 5000
-    i = 2
-    while i < len(sys.argv):
-        if sys.argv[i] == "--port":
-            if i + 1 < len(sys.argv):
-                # Extract the value and remove '--port' and the value from sys.argv
-                port = sys.argv[i + 1]
-                del sys.argv[i : i + 2]
-                break
-        else:
-            i += 1
-
-    # Construct arguments for the flask app minus 'browser' command
-    # plus '--chat'
-    args_plus_chat = ["'{}'".format(s) for s in sys.argv[1:] if s != "browser"] + [
-        '"--chat"'
-    ]
-    formatted_args = ", ".join(args_plus_chat)
-    command = [
-        "flask",
-        "--app",
-        f"{server_dir}/chat_in_browser:create_app(" + formatted_args + ")",
-        "run",
-        "--port",
-        f"{port}",
-    ]
-    subprocess.run(command)
+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)

browser/chat_in_browser.py

@@ -261,12 +261,6 @@ def add_arguments_for_verb(parser, verb: str):
         default=default_model_dir,
         help=f"The directory to store downloaded model artifacts. Default: {default_model_dir}",
     )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=5000,
-        help="Port for the web server in browser mode",
-    )
 
 
 def arg_init(args):