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

# Model adını belirtin
model_name = "MiniMaxAI/MiniMax-M1-40k"

# Tokenizer'ı yükleyin
print(f"Tokenizer yükleniyor: {model_name}...")
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    print("Tokenizer yüklendi.")
except Exception as e:
    print(f"Tokenizer yüklenirken hata oluştu: {e}")
    # Hata durumunda boş bir tokenizer veya hata mesajı
    tokenizer = None

# Modeli yükleyin
print(f"Model yükleniyor: {model_name}...")
# Dikkat: Bu model çok büyük olabilir. device_map="auto" ve torch_dtype=torch.float16
# bellek kullanımını azaltmaya yardımcı olabilir, ancak yeterli değilse CPU'da yüklemek gerekebilir.
try:
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        device_map="auto",
        torch_dtype=torch.float16 # Bellek kullanımını azaltmak ve hesaplamayı hızlandırmak için
        # load_in_8bit=True # 8-bit kuantizasyon kullanmayı deneyebilirsiniz (ek kurulum gerektirir)
    )
    print("Model yüklendi.")
except Exception as e:
    print(f"Model yüklenirken hata oluştu: {e}")
    print("CPU'da yüklemeyi deniyorum...")
    try:
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            device_map="cpu" # Eğer GPU yoksa CPU'da yükle
        )
        print("Model CPU'da yüklendi.")
    except Exception as cpu_e:
        print(f"Model CPU'da yüklenirken de hata oluştu: {cpu_e}")
        model = None

# Modeli değerlendirme moduna (evaluation) alın
if model:
    model.eval()

# Metin üretme fonksiyonu
def generate_text(prompt, max_length=500, temperature=0.7, top_p=0.9, top_k=50):
    """
    Verilen prompt'a göre modelden metin üretir.
    """
    if not model or not tokenizer:
        return "Hata: Model veya tokenizer yüklenemedi."

    # Prompt'u token'lara dönüştür
    try:
        inputs = tokenizer(prompt, return_tensors="pt")
        
        # Eğer model GPU'daysa inputs'ları da GPU'ya taşı
        if model.device != torch.device("cpu"):
            inputs = inputs.to(model.device)
    except Exception as e:
        return f"Prompt tokenleştirilirken hata oluştu: {e}"

    # Metin üretimi
    with torch.no_grad():
        try:
            outputs = model.generate(
                **inputs,
                max_length=max_length,
                do_sample=True,
                temperature=temperature,
                top_p=top_p,
                top_k=top_k,
                pad_token_id=tokenizer.eos_token_id # Padding için EOS token'ını kullan
            )
        except Exception as e:
            return f"Metin üretimi sırasında hata oluştu: {e}"

    # Üretilen token'ları metne dönüştür ve oku
    try:
        generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    except Exception as e:
        return f"Metin çözülürken hata oluştu: {e}"

    return generated_text

# Gradio arayüzünü tanımlayın
with gr.Blocks() as demo:
    gr.Markdown("# MiniMax-M1-40k Demo")
    gr.Markdown("MiniMaxAI/MiniMax-M1-40k modelini kullanarak metin üretin.")

    with gr.Row():
        with gr.Column(scale=3):
            prompt_input = gr.Textbox(label="Prompt", lines=5, placeholder="Buraya metninizi girin...")
            generate_button = gr.Button("Metin Üret")
        with gr.Column(scale=1):
            max_length = gr.Slider(minimum=50, maximum=2048, value=500, step=10, label="Max Uzunluk")
            temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Sıcaklık")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top P")
            top_k = gr.Slider(minimum=1, maximum=100, value=50, step=1, label="Top K")

    output_text = gr.Textbox(label="Üretilen Metin", lines=10, interactive=False)

    # Butona tıklandığında fonksiyonu çağır
    generate_button.click(
        fn=generate_text,
        inputs=[prompt_input, max_length, temperature, top_p, top_k],
        outputs=output_text
    )

# Gradio arayüzünü başlatın
if __name__ == "__main__":
    print("Gradio arayüzü başlatılıyor...")
    demo.launch()