from transformers import pipeline
import gradio as gr

MODELS = {
    "TinyLlama-Chat": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",  # Chat-optimiert!
    "German GPT": "dbmdz/german-gpt2",
    "Mistral-Tiny": "alpindale/Mistral-7B-v0.1-hf"  # Kleinere Mistral-Variante
}

def format_prompt(history, new_message, model_choice):
    """Erstellt modellspezifische Prompt-Formate"""
    if "TinyLlama" in model_choice:
        # TinyLlama benötigt spezielles Format
        prompt = "<|system|>\nDu bist ein hilfsbereiter Assistent.</s>\n"
        for human, assistant in history:
            prompt += f"<|user|>\n{human}</s>\n<|assistant|>\n{assistant}</s>\n"
        prompt += f"<|user|>\n{new_message}</s>\n<|assistant|>\n"
        return prompt
    
    # Allgemeines Format für andere Modelle
    prompt = ""
    for human, assistant in history:
        prompt += f"### Mensch: {human}\n### KI: {assistant}\n"
    prompt += f"### Mensch: {new_message}\n### KI:"
    return prompt

def respond(message, history, model_choice):
    try:
        # Kontext-basierten Prompt erstellen
        full_prompt = format_prompt(history, message, model_choice)
        
        # Pipeline mit optimierten Parametern
        generator = pipeline(
            'text-generation', 
            model=MODELS[model_choice],
            max_new_tokens=150,
            temperature=0.8,  # Mehr Kreativität
            top_p=0.9,
            repetition_penalty=1.2  # Verhindert Wiederholungen
        )
        
        # Antwort generieren
        response = generator(full_prompt)[0]['generated_text']
        
        # Modellspezifische Antwort-Extraktion
        if "TinyLlama" in model_choice:
            return response.split("<|assistant|>")[-1].strip()
        else:
            return response.split("### KI:")[-1].split("###")[0].strip()
        
    except Exception as e:
        return f"❌ Fehler: {str(e)}\n\nTipp: Probier ein anderes Modell oder kürzere Eingaben"

# UI mit besseren Erklärungen
gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Dropdown(
            list(MODELS.keys()), 
            label="Modell", 
            value="TinyLlama-Chat",
            info="TinyLlama-Chat für beste Dialoge"
        )
    ],
    title="🚀 Verbesserter Kontext-Chat",
    description="Wähle 'TinyLlama-Chat' & schreibe ganze Sätze!",
    examples=[
        ["Was ist deine Lieblingsfarbe und warum?"],
        ["Erkläre Quantenphysik in 3 Sätzen"]
    ]
).launch()