From fbec30088c58c704d0309e1231798f09dde8b933 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Mon, 14 Aug 2023 15:49:51 -0700 Subject: [PATCH] v0 litellm --- hf-endpoint/bench.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/hf-endpoint/bench.py b/hf-endpoint/bench.py index be901e3..2e07552 100644 --- a/hf-endpoint/bench.py +++ b/hf-endpoint/bench.py @@ -5,6 +5,7 @@ sys.path.append('../common/') from questions import questions from tqdm import tqdm +from litellm import completion tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b-hf", use_auth_token=True) @@ -30,6 +31,19 @@ def generate(prompt): 'answer': response.json()[0]['generated_text'], 'note': 'hf-endpoint'} +def litellm_generate(prompt, model="meta-llama/llama-2-7b-hf"): + messages=[{"role":"user", "content": prompt}] + start = time.perf_counter() + response = completion(model=model, messages=messages) + request_time = time.perf_counter() - start + text_response = response['choices'][0]['message']['content'] + usage_tokens = response['usage']['completion_tokens'] + return {'tok_count': usage_tokens, + 'time': request_time, + 'question': prompt, + 'answer': text_response, + 'note': 'hf-endpoint'} + if __name__ == '__main__': counter = 1 responses = []