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

def create_pipeline():
    model_name = "m42-health/Llama3-Med42-8B"

    try:
        # Načíst model s podporou pro CPU
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.float32,  # Používáme FP32 pro CPU
            device_map="auto"  # Automatická detekce zařízení (CPU)
        )
    except RuntimeError as e:
        print(f"Chyba při načítání modelu: {e}")
        return None
    
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    pipeline = transformers.pipeline(
        "text-generation",
        model=model,
        tokenizer=tokenizer,
        device_map="auto",  # Použití CPU
    )
    return pipeline

def generate_response(message, history):
    pipeline = create_pipeline()
    
    if pipeline is None:
        return "Model se nepodařilo načíst. Zkuste to prosím později."

    messages = [
        {
            "role": "system",
            "content": (
                "You are a helpful, respectful and honest medical assistant. "
                "You are a second version of Med42 developed by the AI team at M42, UAE. "
                "Always answer in Czech language. "
                "Always answer as helpfully as possible, while being safe. "
            ),
        },
        {"role": "user", "content": message},
    ]
    
    # Předpokládám, že používáte vlastní template pro vytvoření promptu
    prompt = pipeline.tokenizer.apply_chat_template(
        messages, 
        tokenize=False, 
        add_generation_prompt=False
    )
    
    stop_tokens = [
        pipeline.tokenizer.eos_token_id,
        pipeline.tokenizer.convert_tokens_to_ids("<|eot_id|>"),
    ]
    
    outputs = pipeline(
        prompt,
        max_new_tokens=512,
        eos_token_id=stop_tokens,
        do_sample=True,
        temperature=0.4,
        top_k=150,
        top_p=0.75,
    )
    
    response = outputs[0]["generated_text"][len(prompt):]
    return response

# Vytvoření Gradio rozhraní
with gr.Blocks() as demo:
    gr.Markdown("# Med42 Chatbot")
    gr.Markdown("Tento chatbot využívá model Med42-v2-8B pro zodpovídání medicínských dotazů v češtině.")
    
    chatbot = gr.ChatInterface(
        generate_response,
        examples=[
            "Jaké jsou příznaky cukrovky?",
            "Jak se léčí vysoký krevní tlak?",
            "Co je to migréna?",
        ],
        title="Med42 Chatbot",
    )

demo.launch()