# example/transformers-stream/stream.py

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer

model_id = "HIT-SCIR/huozi3.5"

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    attn_implementation="flash_attention_2",
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

tokenizer = AutoTokenizer.from_pretrained(model_id)
tokenizer.chat_template = """{% for message in messages %}{{'<|beginofutterance|>' + message['role'] + '\n' + message['content']}}{% if (loop.last and add_generation_prompt) or not loop.last %}{{ '<|endofutterance|>' + '\n'}}{% endif %}{% endfor %}
{% if add_generation_prompt and messages[-1]['role'] != '助手' %}{{ '<|beginofutterance|>助手\n' }}{% endif %}"""

chat = [
    {"role": "系统", "content": "你是一个智能助手"},
    {"role": "用户", "content": "请你用python写一段快速排序的代码"},
]

inputs = tokenizer.apply_chat_template(
    chat,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
).to(0)

stream_output = model.generate(
    inputs,
    streamer=TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True),
    eos_token_id=57001,
    temperature=0.8,
    top_p=0.9,
    max_new_tokens=2048,
)

# example/vllm-generate/generate.py

from vllm import LLM, SamplingParams

prompts = [
    """<|beginofutterance|>系统
你是一个智能助手<|endofutterance|>
<|beginofutterance|>用户
请你用python写一段快速排序的代码<|endofutterance|>
<|beginofutterance|>助手
""",
]

sampling_params = SamplingParams(
    temperature=0.8, top_p=0.95, stop_token_ids=[57001], max_tokens=2048
)
llm = LLM(
    model="HIT-SCIR/huozi3.5",
    tensor_parallel_size=4,
)
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(generated_text)

# example/openai-api/openai-client.py

from openai import OpenAI

openai_api_key = "EMPTY"
openai_api_base = "https://localhost:8000/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)

chat_response = client.chat.completions.create(
    model="huozi",
    messages=[
        {"role": "系统", "content": "你是一个智能助手"},
        {"role": "用户", "content": "请你用python写一段快速排序的代码"},
    ],
    extra_body={"stop_token_ids": [57001]},
)
print("Chat response:", chat_response.choices[0].message.content)

# example/openai-api/openai-client-gradio.py

from openai import OpenAI
import gradio as gr

openai_api_key = "EMPTY"
openai_api_base = "https://localhost:8000/v1"

client = OpenAI(
    api_key=openai_api_key,
    base_url=openai_api_base,
)


def predict(message, history):
    history_openai_format = [
        {"role": "系统", "content": "你是一个智能助手"},
    ]
    for human, assistant in history:
        history_openai_format.append({"role": "用户", "content": human})
        history_openai_format.append({"role": "助手", "content": assistant})
    history_openai_format.append({"role": "用户", "content": message})
    models = client.models.list()

    stream = client.chat.completions.create(
        model=models.data[0].id,
        messages=history_openai_format,
        temperature=0.8,
        stream=True,
        extra_body={"repetition_penalty": 1, "stop_token_ids": [57001]},
    )

    partial_message = ""
    for chunk in stream:
        partial_message += chunk.choices[0].delta.content or ""
        yield partial_message


gr.ChatInterface(predict).queue().launch()