|
import gradio as gr |
|
from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer |
|
import torch |
|
from threading import Thread |
|
|
|
MODEL_ID = "HODACHI/Llama-3.1-8B-EZO-1.1-it" |
|
DTYPE = torch.bfloat16 |
|
|
|
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID) |
|
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
MODEL_ID, |
|
torch_dtype=torch.bfloat16, |
|
device_map="auto", |
|
low_cpu_mem_usage=True, |
|
) |
|
|
|
pipeline = transformers.pipeline( |
|
"text-generation", |
|
model=model, |
|
tokenizer=tokenizer, |
|
device_map="auto", |
|
) |
|
|
|
def respond( |
|
message, |
|
history: list[tuple[str, str]], |
|
max_tokens, |
|
temperature, |
|
top_p, |
|
): |
|
chat = [] |
|
chat.append({"role": "system", "content": "あなたは誠実で優秀な日本人のアシスタントです。特に指示が無い場合は、原則日本語で回答してください。"}) |
|
for user, assistant in history: |
|
chat.append({"role": "user", "content": user}) |
|
chat.append({"role": "assistant", "content": assistant}) |
|
chat.append({"role": "user", "content": message}) |
|
|
|
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
|
|
|
|
|
streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
thread = Thread(target=model.generate, kwargs=generation_kwargs) |
|
thread.start() |
|
|
|
|
|
|
|
|
|
|
|
outputs = pipeline( |
|
prompt, |
|
max_new_tokens=40, |
|
do_sample=True, |
|
temperature=0.7, |
|
top_p=0.95, |
|
) |
|
|
|
response = outputs[0]["generated_text"] |
|
|
|
demo = gr.ChatInterface( |
|
respond, |
|
additional_inputs=[ |
|
gr.Slider(minimum=1, maximum=2048, value=150, step=1, label="Max new tokens"), |
|
gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"), |
|
gr.Slider( |
|
minimum=0.1, |
|
maximum=1.0, |
|
value=0.95, |
|
step=0.05, |
|
label="Top-p (nucleus sampling)", |
|
), |
|
], |
|
) |
|
|
|
if __name__ == "__main__": |
|
demo.launch() |