Add missing package

Signed-off-by: yuanwu <yuan.wu@intel.com>
This commit is contained in:
yuanwu 2024-10-28 07:04:56 +00:00
parent c23584f626
commit 4c9856f9e5
2 changed files with 4 additions and 2 deletions

View File

@ -14,7 +14,7 @@ from tgi_client import TgiClient
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--server_address", type=str, default="http://localhost:8080", help="Address of the TGI server"
"--server_address", type=str, default="http://localhost:8083", help="Address of the TGI server"
)
parser.add_argument(
"--model_id", type=str, default="meta-llama/Llama-2-7b-chat-hf", help="Model id used in TGI server"
@ -51,7 +51,7 @@ def read_dataset(
)
if len(dataset) > total_sample_count:
dataset = dataset.select(range(total_sample_count))
dataset = dataset.shuffle()
dataset = dataset.shuffle(seed=42)
return [sample["prompt"] for sample in dataset]

View File

@ -1,4 +1,6 @@
import os
import psutil
import signal
import sys
import typer