92 lines
2.3 KiB
Python
92 lines
2.3 KiB
Python
from argparse import ArgumentParser
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
from astrai.inference import InferenceEngine
|
|
from astrai.model import AutoModel
|
|
from astrai.tokenize import AutoTokenizer
|
|
|
|
PROJECT_ROOT = Path(__file__).resolve().parents[2]
|
|
|
|
|
|
def parse_args():
|
|
parser = ArgumentParser(description="Interactive streaming chat")
|
|
parser.add_argument(
|
|
"--model_path",
|
|
type=Path,
|
|
default=PROJECT_ROOT / "params",
|
|
help="Path to model weights (params/ or checkpoint/epoch_N_step_M/)",
|
|
)
|
|
parser.add_argument(
|
|
"--temperature",
|
|
type=float,
|
|
default=0.8,
|
|
help="Sampling temperature (default: 0.8)",
|
|
)
|
|
parser.add_argument(
|
|
"--top_p",
|
|
type=float,
|
|
default=0.95,
|
|
help="Top-p sampling threshold",
|
|
)
|
|
parser.add_argument(
|
|
"--top_k",
|
|
type=int,
|
|
default=50,
|
|
help="Top-k sampling threshold",
|
|
)
|
|
parser.add_argument(
|
|
"--max_tokens",
|
|
type=int,
|
|
default=2048,
|
|
help="Maximum tokens to generate",
|
|
)
|
|
parser.add_argument(
|
|
"--system_prompt",
|
|
type=str,
|
|
default="You are a helpful assistant.",
|
|
help="Optional system prompt",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def chat():
|
|
args = parse_args()
|
|
model_path = args.model_path
|
|
|
|
model = AutoModel.from_pretrained(model_path)
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
model.to(device="cuda", dtype=torch.bfloat16)
|
|
engine = InferenceEngine(model=model, tokenizer=tokenizer)
|
|
|
|
messages = [{"role": "system", "content": args.system_prompt}]
|
|
|
|
while True:
|
|
query = input(">> ")
|
|
if query == "!exit":
|
|
break
|
|
|
|
messages.append({"role": "user", "content": query})
|
|
|
|
full_response = ""
|
|
prompt = tokenizer.apply_chat_template(messages, tokenize=False)
|
|
|
|
for token in engine.generate(
|
|
prompt=prompt,
|
|
stream=True,
|
|
max_tokens=args.max_tokens,
|
|
temperature=args.temperature,
|
|
top_p=args.top_p,
|
|
top_k=args.top_k,
|
|
):
|
|
print(token, end="", flush=True)
|
|
full_response += token
|
|
|
|
print()
|
|
messages.append({"role": "assistant", "content": full_response.strip()})
|
|
|
|
|
|
if __name__ == "__main__":
|
|
chat()
|