Add docstrings for LLM (#137)

cacheflow/entrypoints/llm.py

@@ -11,6 +11,28 @@ from cacheflow.utils import Counter
 
 
 class LLM:
+    """An LLM for generating texts from given prompts and sampling parameters.
+
+    This class includes a tokenizer, a language model (possibly distributed
+    across multiple GPUs), and GPU memory space allocated for intermediate
+    states (aka KV cache). Given a batch of prompts and sampling parameters,
+    this class generates texts from the model, using an intelligent batching
+    mechanism and efficient memory management.
+
+    NOTE: This class is intended to be used for offline inference. For online
+    serving, use the `AsyncLLMServer` class instead.
+    NOTE: For the comprehensive list of arguments, see `ServerArgs`.
+
+    Args:
+        model: The name or path of a HuggingFace Transformers model.
+        tensor_parallel_size: The number of GPUs to use for distributed
+            execution with tensor parallelism.
+        dtype: The data type for the model weights and activations. Currently,
+            we support `float16` and `bfloat16`. If `default`, we use the
+            `torch_dtype` attribute of the model config. If the `torch_dtype`
+            is `float32`, we use `float16` instead.
+        seed: The seed to initialize the random number generator for sampling.
+    """
 
     def __init__(
         self,
@@ -39,19 +61,50 @@ class LLM:
 
     def generate(
         self,
-        prompts: Union[str, List[str]],
+        prompts: Optional[Union[str, List[str]]] = None,
         sampling_params: Optional[SamplingParams] = None,
         prompt_token_ids: Optional[List[List[int]]] = None,
         use_tqdm: bool = True,
     ) -> List[RequestOutput]:
+        """Generates the completions for the input prompts.
+
+        NOTE: This class automatically batches the given prompts, considering
+        the memory constraint. For the best performance, put all of your prompts
+        into a single list and pass it to this method.
+
+        Args:
+            prompts: A list of prompts to generate completions for.
+            sampling_params: The sampling parameters for text generation. If
+                None, we use the default sampling parameters.
+            prompt_token_ids: A list of token IDs for the prompts. If None, we
+                use the tokenizer to convert the prompts to token IDs.
+            use_tqdm: Whether to use tqdm to display the progress bar.
+
+        Returns:
+            A list of `RequestOutput` objects containing the generated
+            completions in the same order as the input prompts.
+        """
+        if prompts is None and prompt_token_ids is None:
+            raise ValueError("Either prompts or prompt_token_ids must be "
+                             "provided.")
         if isinstance(prompts, str):
+            # Convert a single prompt to a list.
             prompts = [prompts]
+        if prompts is not None and prompt_token_ids is not None:
+            if len(prompts) != len(prompt_token_ids):
+                raise ValueError("The lengths of prompts and prompt_token_ids "
+                                 "must be the same.")
         if sampling_params is None:
             # Use default sampling params.
             sampling_params = SamplingParams()
+
         # Add requests to the server.
-        for i in range(len(prompts)):
-            prompt = prompts[i]
+        if prompts is not None:
+            num_requests = len(prompts)
+        else:
+            num_requests = len(prompt_token_ids)
+        for i in range(num_requests):
+            prompt = prompts[i] if prompts is not None else None
             if prompt_token_ids is None:
                 token_ids = None
             else:
@@ -61,7 +114,7 @@ class LLM:
 
     def _add_request(
         self,
-        prompt: str,
+        prompt: Optional[str],
         sampling_params: SamplingParams,
         prompt_token_ids: Optional[List[int]],
     ) -> None: