|
import os |
|
import time |
|
|
|
from openai import OpenAI |
|
from transformers.utils.versions import require_version |
|
|
|
|
|
require_version("openai>=1.5.0", "To fix: pip install openai>=1.5.0") |
|
|
|
|
|
def main(): |
|
client = OpenAI( |
|
api_key="0", |
|
base_url="http://localhost:{}/v1".format(os.environ.get("API_PORT", 8000)), |
|
) |
|
messages = [{"role": "user", "content": "Write a long essay about environment protection as long as possible."}] |
|
num_tokens = 0 |
|
start_time = time.time() |
|
for _ in range(8): |
|
result = client.chat.completions.create(messages=messages, model="test") |
|
num_tokens += result.usage.completion_tokens |
|
|
|
elapsed_time = time.time() - start_time |
|
print("Throughput: {:.2f} tokens/s".format(num_tokens / elapsed_time)) |
|
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
main() |
|
|