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

try:
    # 모델 및 토크나이저 로드
    model_id = "kimhyunwoo/gemma2-ko-dialogue-lora-fp16"
    model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
    # AutoTokenizer 대신 직접 GemmaTokenizer를 로드합니다.
    tokenizer = GemmaTokenizer.from_pretrained(model_id, trust_remote_code=True)

    # CPU 환경에서 모델을 로드하고 사용하는 경우 torch.float32를 명시적으로 사용할 수 있습니다.
    model = model.to(torch.float32)
    model.eval()

    # 텍스트 생성 함수
    def generate_text(text, max_length=50, do_sample=True, temperature=1.0):
        inputs = tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(**inputs, max_length=max_length, do_sample=do_sample, temperature=temperature)
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        return generated_text

    # Gradio 인터페이스 생성
    iface = gr.Interface(
        fn=generate_text,
        inputs=[
            gr.Textbox(lines=5, placeholder="텍스트를 입력하세요..."),
            gr.Slider(minimum=10, maximum=200, step=1, value=50, label="최대 길이"),
            gr.Checkbox(label="샘플링"),
            gr.Slider(minimum=0.1, maximum=2.0, step=0.1, value=1.0, label="온도"),
        ],
        outputs=gr.Textbox(lines=5, label="생성된 텍스트"),
        title="Gemma 2 Text Generator",
        description="Fine-tuned된 Gemma 2 모델을 사용하여 텍스트를 생성합니다.",
    )

    if __name__ == "__main__":
        iface.launch(server_name="0.0.0.0", server_port=7860)
except ImportError as e:
    print(f"ImportError 발생: {e}")
except Exception as e:
    print(f"예상치 못한 오류 발생: {e}")