import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
import torch
from unsloth import FastLanguageModel

# Modell und Tokenizer laden
model_name = "Mario12355/llama_3.1_20.11_fini_dpo"
model, tokenizer = FastLanguageModel.from_pretrained(
    model_name = model_name,
    max_seq_length = 2048,
    dtype = None,
    load_in_4bit = True,
)

# Dein Alpaca-Prompt Template
alpaca_prompt = """Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.

### Instruction:
{}

### Input:
{}

### Response:
{}"""

def translate(
    message,
    history,
    direction,
    max_tokens=128
):
    FastLanguageModel.for_inference(model)
    
    # Richtung bestimmen und Anweisung formatieren
    if direction == "hochdeutsch_to_schwaebisch":
        instruction = "Übersetze den hochdeutschen Text ins Schwäbische. Achte auf eine sinnvolle und korrekte Satzbildung!"
    elif direction == "schwaebisch_to_hochdeutsch":
        instruction = "Übersetze den schwäbischen Text ins Hochdeutsche. Achte auf eine sinnvolle und korrekte Satzbildung!"
    else:
        raise ValueError("Ungültige Übersetzungsrichtung")

    # Input für das Modell vorbereiten
    inputs = tokenizer(
        [alpaca_prompt.format(instruction, message, "")],
        return_tensors="pt"
    ).to(model.device)

    # Streaming-Generator erstellen
    response = ""
    streamer = TextStreamer(tokenizer)
    
    # Generator-Konfiguration
    generation_config = {
        "max_new_tokens": max_tokens,
        "do_sample": True,
        "temperature": 0.7,
        "top_p": 0.95,
        "streamer": streamer,
        **inputs
    }

    # Text generieren und streamen
    for output in model.generate(**generation_config):
        # Token decodieren und zum Response hinzufügen
        new_text = tokenizer.decode(output, skip_special_tokens=True)
        if new_text != response:  # Nur neue Tokens ausgeben
            yield new_text

# Gradio Interface erstellen
demo = gr.ChatInterface(
    translate,
    additional_inputs=[
        gr.Radio(
            choices=["hochdeutsch_to_schwaebisch", "schwaebisch_to_hochdeutsch"],
            value="hochdeutsch_to_schwaebisch",
            label="Übersetzungsrichtung"
        ),
        gr.Slider(
            minimum=32,
            maximum=256,
            value=128,
            step=32,
            label="Maximale Anzahl neuer Tokens"
        )
    ],
    title="Schwäbisch Übersetzer",
    description="""Dieser Übersetzer kann Texte zwischen Hochdeutsch und Schwäbisch übersetzen. 
                   Wählen Sie die gewünschte Übersetzungsrichtung und geben Sie Ihren Text ein.""",
    examples=[
        ["Guten Tag, wie geht es Ihnen?", "hochdeutsch_to_schwaebisch"],
        ["Griaß Gott, wie goht's dir?", "schwaebisch_to_hochdeutsch"]
    ]
)

if __name__ == "__main__":
    demo.launch(
        share=True,
        show_error=True,
    )