import os
import gradio as gr
from transformers import pipeline
import spaces

HF_TOKEN = os.getenv("HF_TOKEN")

# リクエスト時にモデルをロードして使用する関数
@spaces.GPU
def generate_text(
    input_text, 
    max_length=150, 
    temperature=0.7, 
    top_k=50, 
    top_p=0.95, 
    repetition_penalty=1.0, 
    num_return_sequences=1
):
    """ユーザー入力に基づいてテキストを生成し、元のテキストに追加する関数"""
    try:
        # 入力テキストが空の場合は処理しない
        if not input_text.strip():
            return ""
            
        # 関数内でモデルをロードし使用する
        generator = pipeline('text-generation', model='Local-Novel-LLM-project/Vecteus-v1-abliterated', token=HF_TOKEN)
        
        # テキスト生成の設定
        result = generator(
            input_text, 
            max_length=len(input_text.split()) + max_length, 
            temperature=temperature,
            top_k=top_k,
            top_p=top_p,
            repetition_penalty=repetition_penalty,
            num_return_sequences=num_return_sequences,
            pad_token_id=generator.tokenizer.eos_token_id
        )
        
        # 複数系列が生成された場合は、最初の系列を採用し、元のテキストより後の部分のみを抽出
        generated_text = result[0]['generated_text']
        new_text = generated_text[len(input_text):]
        
        # 元のテキストに生成したテキストを追加して返す
        return input_text + new_text
    except Exception as e:
        return f"{input_text}\n\nエラーが発生しました: {str(e)}"

# Gradio インターフェースの作成
with gr.Blocks() as demo:
    gr.Markdown("# テキスト続き生成アシスタント")
    gr.Markdown("テキストボックスに文章を入力し、パラメータを調整して「続きを生成」ボタンをクリックすると、AIがその続きを現在のテキストに追加します。")
    
    # 単一のテキストボックス
    input_text = gr.Textbox(label="テキストを入力してください", placeholder="ここにテキストを入力...", lines=10)
    
    # パラメータの設定UI
    with gr.Row():
        max_length_slider = gr.Slider(minimum=10, maximum=300, value=100, step=10, label="追加する単語数")
        temperature_slider = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="創造性（温度）")
    with gr.Row():
        top_k_slider = gr.Slider(minimum=0, maximum=100, value=50, step=1, label="top_k")
        top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="top_p")
    with gr.Row():
        repetition_penalty_slider = gr.Slider(minimum=1.0, maximum=2.0, value=1.0, step=0.1, label="繰り返しペナルティ")
        num_return_sequences_slider = gr.Slider(minimum=1, maximum=5, value=1, step=1, label="生成する系列数")
    
    # ボタンを横並びに配置
    with gr.Row():
        generate_btn = gr.Button("続きを生成", variant="primary")
        clear_btn = gr.Button("クリア")
    
    # イベントの設定
    generate_btn.click(
        fn=generate_text, 
        inputs=[input_text, max_length_slider, temperature_slider, top_k_slider, top_p_slider, repetition_penalty_slider, num_return_sequences_slider], 
        outputs=input_text
    )
    
    clear_btn.click(lambda: "", None, input_text)
    
    # 使い方の説明
    gr.Markdown("""
    ## 使い方
    1. テキストボックスに続きを生成したい文章を入力します
    2. 生成パラメータ（追加する単語数、創造性、top_k、top_p、繰り返しペナルティ、生成する系列数）を調整します
    3. 「続きを生成」ボタンをクリックすると、入力したテキストの続きが生成され、元のテキストに追加されます
    4. 「クリア」ボタンを押すと、テキストボックスの内容がクリアされます
    5. 満足のいく結果が得られるまで、繰り返し「続きを生成」ボタンを押して文章を発展させることができます
    
    ## ヒント
    - 短い文章から始めて徐々に発展させると良い結果が得られます
    - 創造性（温度）を高くすると予測不可能な生成結果に、低くすると安定した結果になります
    - top_k や top_p、繰り返しペナルティなども状況に応じて調整してみてください
    """)

# アプリの起動
demo.launch()