Separate browser and API, fix context for chat models.

Varun Puri · Varun Puri · commit 301640f6afc0 · 2024-07-08T12:55:41.000-07:00
diff --git a/README.md b/README.md
@@ -54,7 +54,7 @@ source .venv/bin/activate
 ```
 [skip default]: end
 
-[shell default]: ./install_requirements.sh 
+[shell default]: ./install_requirements.sh
 
 Installations can be tested by
 
@@ -118,6 +118,34 @@ python3 torchchat.py generate llama3 --prompt "write me a story about a boy and
 
 For more information run `python3 torchchat.py generate --help`
 
+The `Generator` class can also be imported into a Python program to generate responses.
+
+```
+from generate import Generator, GeneratorArgs
+from build.builder import (
+    BuilderArgs,
+    TokenizerArgs,
+)
+
+...
+
+# Load the model and tokenizer.
+gen = Generator(
+    builder_args,
+    speculative_builder_args,
+    tokenizer_args,
+    generator_args,
+    args.profile,
+    args.quantize,
+    args.draft_quantize,
+)
+
+# The generate function is a Python Generator that will yield torch.Tensors as each token is returned.
+for tok in gen.generate(generator_args):
+  print(gen.tokenizer.decode(y.tolist()))
+
+```
+
 
 ### Browser
 
diff --git a/api/api.py b/api/api.py
@@ -146,17 +146,52 @@ class CompletionResponseChunk():
 class OpenAIAPIGenerator(Generator):
     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 + 1
+            if self.draft_model is not None
+            else self.model.config.max_seq_length
+        )
         
 
     def completion(self, completion_request: CompletionRequest):
+        device_sync(device=self.builder_args.device)
+    
+        id = str(uuid.uuid4())
+        idx = 0
+        buffer = []
+        encoded = self.encode_tokens(
+                        completion_request.prompt, bos=True, device=self.builder_args.device
+                    )
         generator_args = GeneratorArgs(
-            prompt = completion_request.prompt,
-            temperature = completion_request.temperature,
+            completion_request.prompt,
+            encoded_prompt=encoded,
+            chat_mode=False,
         )
-        id = 12345678
-        idx = 0
-        for x, metrics in self.chat(generator_args):
-            content = "".join(self.tokenizer.decode([self.tokenizer.encode(".")[0]] + x.tolist())[1:])
+        
+        def callback(x, *, done_generating=False):
+            return self._callback(
+                x,
+                buffer=buffer,
+                done_generating=done_generating,
+            )
+
+        for y, metrics in self.generate(
+            self.model, 
+            encoded, 
+            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,
+            max_seq_length=self.max_seq_length,
+
+            ):
+            content = "".join(self.tokenizer.decode([self.tokenizer.encode(".")[0]] + y.tolist())[1:])
             chunk_delta = ChunkDelta(
                 role = "assistant",
                 content = content,
@@ -174,6 +209,7 @@ def completion(self, completion_request: CompletionRequest):
                 system_fingerprint = uuid.UUID(int=uuid.getnode()),
             )
             yield chunk_response
+            self.start_pos += y.size(0)
             idx += 1
 
         end_chunk: CompletionChoiceChunk(
@@ -190,10 +226,14 @@ def completion(self, completion_request: CompletionRequest):
             system_fingerprint = uuid.UUID(int=uuid.getnode()),
         )
 
-    def _callback(self, x, *, buffer, period_id, done_generating, tokenizer, is_llama3_model):
-        if x.item() == tokenizer.eos_id():
+    def _callback(self, x, *, buffer, done_generating):
+        period_id = self.tokenizer.encode(".")[0]
+        buffer.append(
+            self.tokenizer.decode([period_id] + x.tolist())[1:]
+        )
+        if x.item() == self.tokenizer.eos_id():
             done_generating = True
-        if is_llama3_model and x.item() == tokenizer.special_tokens["<|eot_id|>"]:
+        if self.is_llama3_model and x.item() == self.tokenizer.special_tokens["<|eot_id|>"]:
             done_generating = True
 
 def main(args):
@@ -248,9 +288,12 @@ def initialize_generator() -> OpenAIAPIGenerator:
             )
 
             def unwrap(completion_generator):
-                for obj in completion_generator:
-                    yield obj.choices[0].delta.content
-
+                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
diff --git a/api/browser.py b/api/browser.py
@@ -0,0 +1,92 @@
+
+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 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("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)
diff --git a/cli.py b/cli.py
@@ -24,7 +24,7 @@
 ).expanduser()
 
 
-KNOWN_VERBS = ["chat", "browser", "download", "generate", "eval", "export", "list", "server", "remove", "where"]
+KNOWN_VERBS = ["chat", "browser", "download", "generate", "eval", "export", "list","remove", "where"]
 
 # Handle CLI arguments that are common to a majority of subcommands.
 def check_args(args, verb: str) -> None:
@@ -256,12 +256,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):
diff --git a/generate.py b/generate.py
@@ -228,8 +228,6 @@ def __init__(
         generator_args.validate_build(self.builder_args)
         generator_args.validate_build(self.speculative_builder_args, "draft model")
 
-        for _ in self.chat(generator_args):
-            continue
 
     def multinomial_sample_one_no_sync(
         self, 
@@ -337,10 +335,11 @@ def decode_n_tokens(
                 input_pos += 1
                 new_tokens.append(next_token.clone())
                 callback(new_tokens[-1], done_generating=_i == num_new_tokens - 2)
-                if need_probs:
+                if need_probs or not next_prob:
+                    yield cur_token.clone(), None
+                else:
                     new_probs.append(next_prob.clone())
-
-                yield cur_token.clone(), next_prob.clone()
+                    yield cur_token.clone(), next_prob.clone()
                 cur_token = next_token
 
                 # encountered eos
@@ -365,7 +364,7 @@ def decode_n_tokens(
                 model, eos_token.view(1, -1), input_pos, need_probs, **sampling_kwargs
             )
             input_pos += 1
-            yield eos_token.clone(), next_prob.clone()
+            yield eos_token.clone(), (next_prob.clone() if next_prob else None)
 
         # return new_tokens, new_probs
 
@@ -450,9 +449,7 @@ def generate(
         speculate_k: Optional[int] = 8,
         sequential_prefill=True,
         callback=lambda x: x,
-        tokenizer=None,
         max_seq_length: int,
-        is_llama3_model: bool = False,
         **sampling_kwargs,
     ) -> torch.Tensor:
         """
@@ -536,8 +533,8 @@ def generate(
                 max_new_tokens - 1,
                 callback=callback,
                 need_probs=False,
-                eos_token_id=tokenizer.eos_id() if tokenizer else 2,
-                eot_id=tokenizer.special_tokens["<|eot_id|>"] if is_llama3_model else None,
+                eos_token_id=self.tokenizer.eos_id() if self.tokenizer else 2,
+                eot_id=self.tokenizer.special_tokens["<|eot_id|>"] if self.is_llama3_model else None,
                 **sampling_kwargs,
             ):
                 generated_tokens.append(generated_token)
@@ -555,19 +552,20 @@ def generate(
         return seq, generate_stats
 
 
-    def encode_tokens(self, tokenizer, string, bos=True, device="cpu"):
+    def encode_tokens(self, string, bos=True, device="cpu"):
         tokens = self.tokenizer.encode(string)
         if bos:
-            tokens = [tokenizer.bos_id()] + tokens
+            tokens = [self.tokenizer.bos_id()] + tokens
         return torch.tensor(tokens, dtype=torch.int, device=device)
 
-    def _callback(self, x, *, buffer, period_id, done_generating, tokenizer, is_llama3_model):
+    def _callback(self, x, *, buffer, done_generating):
+        period_id = self.tokenizer.encode(".")[0]
         buffer.append(
             self.tokenizer.decode([period_id] + x.tolist())[1:]
         )  # I think this results in the first output token being dropped from the display which is wrong.
         if x.item() == self.tokenizer.eos_id():
             done_generating = True
-        if is_llama3_model and x.item() == self.tokenizer.special_tokens["<|eot_id|>"]:
+        if self.is_llama3_model and x.item() == self.tokenizer.special_tokens["<|eot_id|>"]:
             done_generating = True
             buffer = buffer[:-1]  # drop the eot_id from the output buffer
         if len(buffer) == 4 or done_generating:
@@ -581,7 +579,7 @@ def chat(
     ):
         print("Starting Interactive Chat")
         encoded = self.encode_tokens(
-            self.tokenizer, generator_args.prompt, bos=True, device=self.builder_args.device
+            generator_args.prompt, bos=True, device=self.builder_args.device
         )
         logging.debug(encoded)
         prompt_length = encoded.size(0)
@@ -664,7 +662,7 @@ def chat(
                     else:
                         prompt = f"{B_INST} {prompt.strip()} {E_INST}"
                     encoded = self.encode_tokens(
-                        self.tokenizer, prompt, bos=True, device=self.builder_args.device
+                        prompt, bos=True, device=self.builder_args.device
                     )
                 else:
                     if self.system_prompt is not None:
@@ -707,10 +705,7 @@ def callback(x, *, done_generating=False):
                     return self._callback(
                         x,
                         buffer=buffer,
-                        period_id=period_id,
                         done_generating=done_generating,
-                        tokenizer=self.tokenizer,
-                        is_llama3_model=self.is_llama3_model,
                     )
 
             else:
@@ -722,10 +717,7 @@ def callback(x, *, done_generating=False):
                     return self._callback(
                         x,
                         buffer=buffer,
-                        period_id=period_id,
                         done_generating=done_generating,
-                        tokenizer=self.tokenizer,
-                        is_llama3_model=self.is_llama3_model,
                     )
 
             if (i != generator_args.num_samples - 1 or not self.profile) or (
@@ -751,9 +743,7 @@ def callback(x, *, done_generating=False):
                     top_k=generator_args.top_k,
                     sequential_prefill=generator_args.sequential_prefill,
                     start_pos=start_pos,
-                    tokenizer=self.tokenizer,
                     max_seq_length=max_seq_length,
-                    is_llama3_model=self.is_llama3_model,
                 ):
                     if metrics:
                         aggregate_metrics["accept_counts"].append(metrics["accept_counts"])
diff --git a/torchchat.py b/torchchat.py
