Create predict.py

predict.py

@@ -0,0 +1,148 @@
+import os
+import time
+
+os.environ["HF_HUB_ENABLE_HF_TRANSFER"] = "1"
+
+from vllm import LLM, SamplingParams
+import torch
+from cog import BasePredictor, Input, ConcatenateIterator
+import typing as t
+
+
+MODEL_ID = "TheBloke/Mistral-7B-OpenOrca-AWQ"
+PROMPT_TEMPLATE = """\
+<|im_start|>system
+You are MistralOrca, a large language model trained by Alignment Lab AI. Write out your reasoning step-by-step to be sure you get the right answers!
+<|im_end|>
+<|im_start|>user
+{prompt}<|im_end|>
+<|im_start|>assistant
+"""
+
+DEFAULT_MAX_NEW_TOKENS = 512
+DEFAULT_TEMPERATURE = 0.8
+DEFAULT_TOP_P = 0.95
+DEFAULT_TOP_K = 50
+DEFAULT_PRESENCE_PENALTY = 0.0  # 1.15
+DEFAULT_FREQUENCY_PENALTY = 0.0  # 0.2
+
+
+def vllm_generate_iterator(
+    self, prompt: str, /, *, echo: bool = False, stop: str  = None, stop_token_ids: t.List[int] = None, sampling_params=None, **attrs: t.Any
+) -> t.Iterator[t.Dict[str, t.Any]]:
+    request_id: str = attrs.pop('request_id', None)
+    if request_id is None: raise ValueError('request_id must not be None.')
+    if stop_token_ids is None: stop_token_ids = []
+    stop_token_ids.append(self.tokenizer.eos_token_id)
+    stop_ = set()
+    if isinstance(stop, str) and stop != '': stop_.add(stop)
+    elif isinstance(stop, list) and stop != []: stop_.update(stop)
+    for tid in stop_token_ids:
+        if tid: stop_.add(self.tokenizer.decode(tid))
+
+    # if self.config['temperature'] <= 1e-5: top_p = 1.0
+    # else: top_p = self.config['top_p']
+    # config = self.config.model_construct_env(stop=list(stop_), top_p=top_p, **attrs)
+    self.add_request(request_id=request_id, prompt=prompt, sampling_params=sampling_params)
+
+    token_cache = []
+    print_len = 0
+
+    while self.has_unfinished_requests():
+        for request_output in self.step():
+            # Add the new tokens to the cache
+            for output in request_output.outputs:
+                text = output.text
+                yield {'text': text, 'error_code': 0, 'num_tokens': len(output.token_ids)}
+
+            if request_output.finished: break
+
+
+class Predictor(BasePredictor):
+
+    def setup(self):
+        self.llm = LLM(
+            model=MODEL_ID,
+            quantization="awq",
+            dtype="float16"
+        )
+
+    def predict(
+        self,
+        prompt: str,
+        max_new_tokens: int = Input(
+            description="The maximum number of tokens the model should generate as output.",
+            default=DEFAULT_MAX_NEW_TOKENS,
+        ),
+        temperature: float = Input(
+            description="The value used to modulate the next token probabilities.", default=DEFAULT_TEMPERATURE
+        ),
+        top_p: float = Input(
+            description="A probability threshold for generating the output. If < 1.0, only keep the top tokens with cumulative probability >= top_p (nucleus filtering). Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751).",
+            default=DEFAULT_TOP_P,
+        ),
+        top_k: int = Input(
+            description="The number of highest probability tokens to consider for generating the output. If > 0, only keep the top k tokens with highest probability (top-k filtering).",
+            default=DEFAULT_TOP_K,
+        ),
+        presence_penalty: float = Input(
+            description="Presence penalty",
+            default=DEFAULT_PRESENCE_PENALTY,
+        ),
+        frequency_penalty: float = Input(
+            description="Frequency penalty",
+            default=DEFAULT_FREQUENCY_PENALTY,
+        ),
+        prompt_template: str = Input(
+            description="The template used to format the prompt. The input prompt is inserted into the template using the `{prompt}` placeholder.",
+            default=PROMPT_TEMPLATE,
+        )
+    ) -> ConcatenateIterator:
+        prompts = [
+            (
+                prompt_template.format(prompt=prompt),
+                SamplingParams(
+                    max_tokens=max_new_tokens,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    presence_penalty=presence_penalty,
+                    frequency_penalty=frequency_penalty
+                )
+            )
+        ]
+        start = time.time()
+        while True:
+            if prompts:
+                prompt, sampling_params = prompts.pop(0)
+                gen = vllm_generate_iterator(self.llm.llm_engine, prompt, echo=False, stop=None, stop_token_ids=None, sampling_params=sampling_params, request_id=0)
+                last = ""
+                for _, x in enumerate(gen):
+                    if x['text'] == "":
+                        continue
+                    yield x['text'][len(last):]
+                    last = x["text"]
+                    num_tokens = x["num_tokens"]
+                print(f"\nGenerated {num_tokens} tokens in {time.time() - start} seconds.")
+
+                if not (self.llm.llm_engine.has_unfinished_requests() or prompts):
+                    break
+
+
+if __name__ == '__main__':
+    import sys
+    p = Predictor()
+    p.setup()
+    gen = p.predict(
+        "Write me an itinerary for my dog's birthday party.",
+        512,
+        0.8,
+        0.95,
+        50,
+        1.0,
+        0.2,
+        PROMPT_TEMPLATE,
+    )
+    for out in gen:
+        print(out, end="")
+        sys.stdout.flush()