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

# Qwenモデルの読み込み
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")

# モデルがテキストを生成する関数
def generate_text(prompt, max_length, temperature, top_p):
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(
        **inputs,
        max_length=max_length,
        temperature=temperature,
        top_p=top_p,
        do_sample=True
    )
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return generated_text

# テキストを音声に変換する関数
def text_to_speech(text):
    tts = gTTS(text=text, lang='ja')
    file_path = "output.mp3"
    tts.save(file_path)
    return file_path

# モデルの出力と音声変換を組み合わせた関数
def generate_and_speak(prompt, max_length, temperature, top_p):
    generated_text = generate_text(prompt, max_length, temperature, top_p)
    speech_file = text_to_speech(generated_text)
    return generated_text, speech_file

# Gradioインターフェースの作成
with gr.Blocks() as demo:
    # 入力欄（プロンプト、長さ、温度、Top-p）
    prompt = gr.Textbox(label="プロンプトを入力してください")
    max_length = gr.Slider(10, 1000, value=400, step=10, label="最大長")
    temperature = gr.Slider(0.1, 1.0, value=0.7, step=0.1, label="生成温度")
    top_p = gr.Slider(0.0, 1.0, value=0.95, step=0.05, label="Top-p")

    # 出力欄（生成されたテキストと音声ファイル）
    generated_text_output = gr.Textbox(label="生成されたテキスト")
    audio_output = gr.Audio(label="生成された音声")
    
    # ボタン（生成ボタンとクリアボタン）
    generate_button = gr.Button("生成")
    clear_button = gr.Button("クリア")

    # ボタンが押されたときのアクション
    generate_button.click(generate_and_speak, 
                          inputs=[prompt, max_length, temperature, top_p], 
                          outputs=[generated_text_output, audio_output])
    
    # クリアボタンの機能
    clear_button.click(lambda: ("", None), None, [prompt, generated_text_output, audio_output])

# インターフェースの起動
demo.launch()