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

# Modelle und Tokenizer laden
# model_names = {
#     "LeoLM_13B": "LeoLM/leo-hessianai-13b",
#     "Occiglot_7B": "occiglot/occiglot-7b-de-en",
#     "LLaMA2_13B": "meta-llama/Llama-2-13b-hf"
# }

model_names = {
    "LeoLM_7B": "LeoLM/leo-hessianai-7b",
    "Occiglot_7B": "occiglot/occiglot-7b-de-en"
}

# Tokenizer und Modelle vorbereiten
tokenizers = {name: AutoTokenizer.from_pretrained(model) for name, model in model_names.items()}
models = {name: AutoModelForCausalLM.from_pretrained(model) for name, model in model_names.items()}

# Funktion zur Generierung der Antwort
def generate_response(model_choice, prompt):
    tokenizer = tokenizers[model_choice]
    model = models[model_choice]
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(inputs["input_ids"], max_new_tokens=100, do_sample=True)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# Vergleich von LLMs: LeoLM und Occiglot")
    with gr.Row():
        model_choice = gr.Radio(list(model_names.keys()), label="Modell auswählen")
        prompt = gr.Textbox(label="Frage stellen", placeholder="Was sind die Hauptursachen für Bluthochdruck?")
        output = gr.Textbox(label="Antwort")

    submit_button = gr.Button("Antwort generieren")
    submit_button.click(generate_response, inputs=[model_choice, prompt], outputs=output)

# App starten
demo.launch()