Merge branch 'main' into benchmarking_script

Author: Jack-Khuu

README.md

@@ -181,16 +181,6 @@ This mode generates text based on an input prompt.
 python3 torchchat.py generate llama3.1 --prompt "write me a story about a boy and his bear"
 ```
 
-### Browser
-This mode allows you to chat with the model using a UI in your browser
-Running the command automatically open a tab in your browser.
-
-[skip default]: begin
-
-```
-streamlit run torchchat.py -- browser llama3.1
-```
-
 [skip default]: end
 
 ### Server
@@ -252,6 +242,19 @@ curl http://127.0.0.1:5000/v1/chat \
 
 </details>
 
+### Browser
+This command opens a basic browser interface for local chat by querying a local server.
+
+First, follow the steps in the Server section above to start a local server. Then, in another terminal, launch the interface. Running the following will open a tab in your browser.
+
+[skip default]: begin
+
+```
+streamlit run browser/browser.py
+```
+
+Use the "Max Response Tokens" slider to limit the maximum number of tokens generated by the model for each response. Click the "Reset Chat" button to remove the message history and start a fresh chat.
+
 
 ## Desktop/Server Execution
 

api/api.py

@@ -10,6 +10,8 @@
 from dataclasses import dataclass
 from typing import Any, Dict, List, Optional, Union
 
+import torch
+
 from build.utils import device_sync
 
 from generate import Generator, GeneratorArgs
@@ -123,6 +125,9 @@ class CompletionRequest:
     parallel_tool_calls: Optional[bool] = None  # unimplemented - Assistant features
     user: Optional[str] = None  # unimplemented
 
+    def __post_init__(self):
+        self.stream = bool(self.stream)
+
 
 @dataclass
 class CompletionChoice:
@@ -202,7 +207,7 @@ class CompletionResponseChunk:
     choices: List[CompletionChoiceChunk]
     created: int
     model: str
-    system_fingerprint: str
+    system_fingerprint: Optional[str] = None
     service_tier: Optional[str] = None
     object: str = "chat.completion.chunk"
     usage: Optional[UsageStats] = None
@@ -222,7 +227,6 @@ def __init__(self, *args, **kwargs):
         """
 
         super().__init__(*args, **kwargs)
-        self.start_pos = 0
         self.max_seq_length = (
             self.model.config.max_seq_length
             + self.speculative_builder_args.speculate_k
@@ -257,20 +261,25 @@ def chunked_completion(self, completion_request: CompletionRequest):
             CompletionResponseChunk objects in response to completion_request as tokens are generated.
 
         """
-        device_sync(device=self.builder_args.device)
 
         # Initialize counters for chunk responses and encode the prompt.
         id = str(uuid.uuid4())
 
         idx = 0
-        buffer = []
-        encoded = self.encode_tokens(
-            completion_request.messages[-1].get("content"),
-            bos=True,
-            device=self.builder_args.device,
+        tokens = self.chat_formatter.encode_dialog_prompt(
+            dialog=[
+                {"role": message["role"], "content": message["content"]}
+                for message in completion_request.messages
+            ]
         )
+
+        encoded = torch.tensor(tokens, dtype=torch.int, device=self.builder_args.device)
+        print(self.tokenizer.decode(tokens))
+
+        start_pos = 0
+
         generator_args = GeneratorArgs(
-            completion_request.messages[-1].get("content"),
+            None,
             max_new_tokens=(
                 int(completion_request.max_tokens)
                 if completion_request.max_tokens
@@ -279,33 +288,39 @@ def chunked_completion(self, completion_request: CompletionRequest):
             encoded_prompt=encoded,
             temperature=float(completion_request.temperature),
             chat_mode=False,
+            sequential_prefill=True,
         )
 
         def callback(x, *, done_generating=False):
             return self._callback(
                 x,
-                buffer=buffer,
+                buffer=None,
                 done_generating=done_generating,
             )
 
+        device_sync(device=self.builder_args.device)
+
         # Process each token, metrics tuple yielded by Generator.generate.
         for y, _ in self.generate(
-            self.model,
-            encoded,
-            generator_args.max_new_tokens,
+            model=self.model,
+            prompt=encoded,
+            max_new_tokens=generator_args.max_new_tokens,
             draft_model=self.draft_model,
             speculate_k=generator_args.speculate_k,
             chat_mode=generator_args.chat_mode,
             callback=callback,
             temperature=generator_args.temperature,
             top_k=generator_args.top_k,
             sequential_prefill=generator_args.sequential_prefill,
-            start_pos=self.start_pos,
+            start_pos=start_pos,
             max_seq_length=self.max_seq_length,
-            seed=int(completion_request.seed),
+            seed=int(completion_request.seed or 0),
         ):
             if y is None:
                 continue
+            elif y.item() == self.tokenizer.eos_id:
+                # Stop generation if the EOS token is generated.
+                break
 
             # Decode the torch.Tensor token to a string and append to the buffer. Separate the sequences with a period token.
             content = "".join(
@@ -321,16 +336,17 @@ def callback(x, *, done_generating=False):
             choice_chunk = CompletionChoiceChunk(
                 delta=chunk_delta,
                 index=idx,
+                finish_reason=None,
             )
             chunk_response = CompletionResponseChunk(
-                id=str(id),
+                id="chatcmpl-" + str(id),
                 choices=[choice_chunk],
                 created=int(time.time()),
                 model=completion_request.model,
                 system_fingerprint=self.system_fingerprint,
             )
             yield chunk_response
-            self.start_pos += y.size(0)
+            start_pos += y.size(0)
             idx += 1
 
         # Yield an ending chunk indicating the generation has completed.
@@ -339,7 +355,7 @@ def callback(x, *, done_generating=False):
         )
 
         yield CompletionResponseChunk(
-            id=str(id),
+            id="chatcmpl-" + str(id),
             choices=[end_chunk],
             created=int(time.time()),
             model=completion_request.model,
@@ -355,7 +371,7 @@ def sync_completion(self, request: CompletionRequest):
 
         message = AssistantMessage(content=output)
         return CompletionResponse(
-            id=str(uuid.uuid4()),
+            id="chatcmpl-" + str(uuid.uuid4()),
             choices=[
                 CompletionChoice(
                     finish_reason="stop",
@@ -369,10 +385,4 @@ def sync_completion(self, request: CompletionRequest):
         )
 
     def _callback(self, x, *, buffer, done_generating):
-        period_id = self.tokenizer.encode(".")[0]
-        buffer.append(self.tokenizer.decode([period_id] + x.tolist())[1:])
-        if (
-            self.is_llama3_model
-            and x.item() == self.tokenizer.special_tokens["<|eot_id|>"]
-        ):
-            buffer = buffer[:-1]  # drop the eot_id from the output buffer
+        pass

browser/browser.py

@@ -1,91 +1,67 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-
-# This source code is licensed under the license found in the
-# LICENSE file in the root directory of this source tree.
-
 import time
-
 import streamlit as st
-from api.api import CompletionRequest, OpenAiApiGenerator
-
-from build.builder import BuilderArgs, TokenizerArgs
-
-from generate import GeneratorArgs
-
-
-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()
-
-    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=[],
+from openai import OpenAI
+
+st.title("torchchat")
+
+start_state = [
+    {
+        "role": "system",
+        "content": "You're an assistant. Answer questions directly, be brief, and have fun.",
+    },
+    {"role": "assistant", "content": "How can I help you?"},
+]
+
+with st.sidebar:
+    response_max_tokens = st.slider(
+        "Max Response Tokens", min_value=10, max_value=1000, value=250, step=10
+    )
+    if st.button("Reset Chat", type="primary"):
+        st.session_state["messages"] = start_state
+
+if "messages" not in st.session_state:
+    st.session_state["messages"] = start_state
+
+
+for msg in st.session_state.messages:
+    st.chat_message(msg["role"]).write(msg["content"])
+
+if prompt := st.chat_input():
+    client = OpenAI(
+        base_url="http://127.0.0.1:5000/v1",
+        api_key="813",  # The OpenAI API requires an API key, but since we don't consume it, this can be any non-empty string.
+    )
+
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    st.chat_message("user").write(prompt)
+
+    with st.chat_message("assistant"), st.status(
+        "Generating... ", expanded=True
+    ) as status:
+
+        def get_streamed_completion(completion_generator):
+            start = time.time()
+            tokcount = 0
+            for chunk in completion_generator:
+                tokcount += 1
+                yield chunk.choices[0].delta.content
+
+            status.update(
+                label="Done, averaged {:.2f} tokens/second".format(
+                    tokcount / (time.time() - start)
+                ),
+                state="complete",
             )
 
-            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(
+            get_streamed_completion(
+                client.chat.completions.create(
+                    model="llama3",
+                    messages=st.session_state.messages,
+                    max_tokens=response_max_tokens,
+                    stream=True,
                 )
+            )
+        )[0]
 
-            response = st.write_stream(unwrap(gen.completion(req)))
-
-        # Add assistant response to chat history
-        st.session_state.messages.append({"role": "assistant", "content": response})
+    st.session_state.messages.append({"role": "assistant", "content": response})

build/builder.py

@@ -400,7 +400,7 @@ def _maybe_parellelize_model(
     if the user specifies using distributed inference. If not, this is a no-op.
 
     Args:
-        module (:class:`nn.Module`):
+        model (:class:`nn.Module`):
             Module to be parallelized.
         builder_args (:class:`BuilderArgs`):
             Command args for model building.
@@ -440,6 +440,7 @@ def _initialize_model(
     quantize,
     tokenizer=None,
     max_seq_length=None,
+    support_tensor_subclass: bool = True,
 ):
     print("Loading model...")
 
@@ -510,7 +511,13 @@ def _initialize_model(
         if quantize:
             print(f"Quantizing the model with: {quantize}")
             with measure_time("Time to quantize model: {time:.02f} seconds"):
-                quantize_model(model, builder_args.device, quantize, tokenizer)
+                quantize_model(
+                    model,
+                    builder_args.device,
+                    quantize,
+                    tokenizer,
+                    support_tensor_subclass,
+                )
                 device_sync(device=builder_args.device)
 
         if builder_args.setup_caches: