import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

model_id = "llm-jp/llm-jp-13b-instruct-full-dolly-ichikara_004_001_single-oasst-oasst2-v2.0"
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32

tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=dtype, device_map="auto", token=token)


def generate_response(prompt):
    chat = [
        {"role": "system", "content": "以下は、タスクを説明する指示です。要求を適切に満たす応答を書きなさい。"},
        {"role": "user", "content": prompt},
    ]
    tokenized_input = tokenizer.apply_chat_template(
        chat, add_generation_prompt=True, tokenize=True, return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():
        output = model.generate(
            tokenized_input,
            max_new_tokens=100,
            do_sample=True,
            top_p=0.95,
            temperature=0.7,
            repetition_penalty=1.05,
        )[0]

    return tokenizer.decode(output)


@spaces.GPU
def process_input(prompt):
    response = generate_response(prompt)
    return response


demo = gr.Interface(
    fn=process_input,
    inputs=gr.Textbox(lines=3, placeholder="ここに質問を入力してください..."),
    outputs=gr.Textbox(),
    title="LLM-JP 対話デモ",
    description="このデモはLLM-JPモデルを使用して、与えられた質問に対する応答を生成します。",
    examples=[
        ["自然言語処理とは何か"],
        ["機械学習の基本的な概念を説明してください"],
        ["プログラミング言語Pythonの特徴は何ですか"],
        ["人工知能（AI）の倫理的な課題について教えてください"],
    ],
)

if __name__ == "__main__":
    demo.launch()