| import argparse | |
| import json | |
| import requests | |
| from llava.conversation import default_conversation | |
| def main(): | |
| if args.worker_address: | |
| worker_addr = args.worker_address | |
| else: | |
| controller_addr = args.controller_address | |
| ret = requests.post(controller_addr + "/refresh_all_workers") | |
| ret = requests.post(controller_addr + "/list_models") | |
| models = ret.json()["models"] | |
| models.sort() | |
| print(f"Models: {models}") | |
| ret = requests.post(controller_addr + "/get_worker_address", | |
| json={"model": args.model_name}) | |
| worker_addr = ret.json()["address"] | |
| print(f"worker_addr: {worker_addr}") | |
| if worker_addr == "": | |
| return | |
| conv = default_conversation.copy() | |
| conv.append_message(conv.roles[0], args.message) | |
| prompt = conv.get_prompt() | |
| headers = {"User-Agent": "LLaVA Client"} | |
| pload = { | |
| "model": args.model_name, | |
| "prompt": prompt, | |
| "max_new_tokens": args.max_new_tokens, | |
| "temperature": 0.7, | |
| "stop": conv.sep, | |
| } | |
| response = requests.post(worker_addr + "/worker_generate_stream", headers=headers, | |
| json=pload, stream=True) | |
| print(prompt.replace(conv.sep, "\n"), end="") | |
| for chunk in response.iter_lines(chunk_size=8192, decode_unicode=False, delimiter=b"\0"): | |
| if chunk: | |
| data = json.loads(chunk.decode("utf-8")) | |
| output = data["text"].split(conv.sep)[-1] | |
| print(output, end="\r") | |
| print("") | |
| if __name__ == "__main__": | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--controller-address", type=str, default="http://localhost:21001") | |
| parser.add_argument("--worker-address", type=str) | |
| parser.add_argument("--model-name", type=str, default="facebook/opt-350m") | |
| parser.add_argument("--max-new-tokens", type=int, default=32) | |
| parser.add_argument("--message", type=str, default= | |
| "Tell me a story with more than 1000 words.") | |
| args = parser.parse_args() | |
| main() | |