import gradio as gr
from peft import PeftModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import time

# モデルとトークナイザーのロード
base_model = AutoModelForCausalLM.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct")
tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/SmolLM2-135M-Instruct")

# LoRA（またはPEFT）アダプターを適用
model = PeftModel.from_pretrained(base_model, "summerstars/beachball1.00")

# 汎用的なプロンプト（モデル名を「summer」と答える）
system_prompt = (
    "You are an AI assistant. Answer the following questions clearly and concisely."
    " For any question asking about the model name, always respond with 'The model name is summer.'"
)

# 出力のタイピング風生成
def generate_response(input_text):
    # プロンプトにユーザーの入力を追加
    full_prompt = system_prompt + "\nQuestion: " + input_text + "\nAnswer:"

    # トークナイズ
    inputs = tokenizer(full_prompt, return_tensors="pt")

    # モデルを使って応答を生成
    output = model.generate(inputs["input_ids"], max_length=200, num_return_sequences=1)

    # 出力されたトークンをデコードして生成されたテキストを取得
    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    
    # 文字列をタイピングのように逐次的に表示
    response = ""
    for char in generated_text:
        response += char
        time.sleep(0.05)  # 文字を1つずつ表示するための遅延
        yield response  # 一時停止して応答を逐次的に返す

# Gradioインターフェースの設定
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="質問を入力してください...", label="質問"),
    outputs=gr.Textbox(lines=10, interactive=False, label="AIの応答"),
    live=True,  # ユーザーが入力中にリアルタイムで反応する
    theme="huggingface",  # UIテーマの設定（他にも色々なテーマが選べる）
)

# インターフェースを実行
iface.launch()