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()