From 4c9856f9e59ea0f3b3436a87ec3c7cb6394fd9ee Mon Sep 17 00:00:00 2001 From: yuanwu Date: Mon, 28 Oct 2024 07:04:56 +0000 Subject: [PATCH] Add missing package Signed-off-by: yuanwu --- examples/run_generation.py | 4 ++-- server/text_generation_server/cli.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/run_generation.py b/examples/run_generation.py index 81423c38..fedfc259 100644 --- a/examples/run_generation.py +++ b/examples/run_generation.py @@ -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] diff --git a/server/text_generation_server/cli.py b/server/text_generation_server/cli.py index 756322c9..f9b4caa9 100644 --- a/server/text_generation_server/cli.py +++ b/server/text_generation_server/cli.py @@ -1,4 +1,6 @@ import os +import psutil +import signal import sys import typer