File size: 3,727 Bytes
90fad3d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | import argparse
import json
import os
from openai import OpenAI
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Tool for checking if model deployment is successful",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Usage examples:
python scripts/check_deployment_en.py --base-url http://localhost:8000/v1 --apikey your-key --model autoglm-phone-9b
python scripts/check_deployment_en.py --base-url http://localhost:8000/v1 --apikey your-key --model autoglm-phone-9b --messages-file custom.json
""",
)
parser.add_argument(
"--base-url",
type=str,
required=True,
help="Base URL of the API service, e.g.: http://localhost:8000/v1",
)
parser.add_argument(
"--apikey", type=str, default="EMPTY", help="API key (default: EMPTY)"
)
parser.add_argument(
"--model",
type=str,
required=True,
help="Name of the model to test, e.g.: autoglm-phone-9b",
)
parser.add_argument(
"--messages-file",
type=str,
default="scripts/sample_messages_en.json",
help="Path to JSON file containing test messages (default: scripts/sample_messages_en.json)",
)
parser.add_argument(
"--max-tokens",
type=int,
default=3000,
help="Maximum generation tokens (default: 3000)",
)
parser.add_argument(
"--temperature",
type=float,
default=0.0,
help="Sampling temperature (default: 0.0)",
)
parser.add_argument(
"--top_p",
type=float,
default=0.85,
help="Nucleus sampling parameter (default: 0.85)",
)
parser.add_argument(
"--frequency_penalty",
type=float,
default=0.2,
help="Frequency penalty parameter (default: 0.2)",
)
args = parser.parse_args()
# Read test messages
if not os.path.exists(args.messages_file):
print(f"Error: Message file {args.messages_file} does not exist")
exit(1)
with open(args.messages_file) as f:
messages = json.load(f)
base_url = args.base_url
api_key = args.apikey
model = args.model
print(f"Starting model inference test...")
print(f"Base URL: {base_url}")
print(f"Model: {model}")
print(f"Messages file: {args.messages_file}")
print("=" * 80)
try:
client = OpenAI(
base_url=base_url,
api_key=api_key,
)
response = client.chat.completions.create(
messages=messages,
model=model,
max_tokens=args.max_tokens,
temperature=args.temperature,
top_p=args.top_p,
frequency_penalty=args.frequency_penalty,
stream=False,
)
print("\nModel inference result:")
print("=" * 80)
print(response.choices[0].message.content)
print("=" * 80)
if response.usage:
print(f"\nStatistics:")
print(f" - Prompt tokens: {response.usage.prompt_tokens}")
print(f" - Completion tokens: {response.usage.completion_tokens}")
print(f" - Total tokens: {response.usage.total_tokens}")
print(
f"\nPlease evaluate the above inference result to determine if the model deployment meets expectations."
)
except Exception as e:
print(f"\nError occurred while calling API:")
print(f"Error type: {type(e).__name__}")
print(f"Error message: {str(e)}")
print(
"\nTip: Please check if base_url, api_key and model parameters are correct, and if the service is running."
)
exit(1)
|