import gradio as gr
from transformers import pipeline
import torch

# Titel und Beschreibung
title = "CocoAi-1PT Demo"
description = "Dies ist eine Demo der KI CocoAi-1PT von CocoEntertainment."

# 1. Lade das Modell
# WICHTIG: Wir nutzen torch_dtype=torch.float16, um den RAM-Verbrauch zu halbieren!
# Sonst stürzt der Space bei großen Modellen auf der Free-Tier ab.
try:
    generator = pipeline(
        "text-generation", 
        model="CocoEntertainment/CocoAi-1PT",
        torch_dtype=torch.float16, 
        device_map="auto" # Versucht, CPU/RAM optimal zu nutzen
    )
except Exception as e:
    generator = None
    print(f"Fehler beim Laden des Modells: {e}")

# 2. Definiere die Funktion
def generate_text(prompt):
    if generator is None:
        return "Fehler: Das Modell konnte nicht geladen werden (evtl. zu groß für Free Tier)."
    
    try:
        # Generierungsparameter
        results = generator(
            prompt, 
            max_new_tokens=100,  # Begrenzt die Antwortlänge (spart Zeit)
            do_sample=True, 
            temperature=0.7,
            truncation=True
        )
        return results[0]['generated_text']
    except Exception as e:
        return f"Ein Fehler ist aufgetreten: {str(e)}"

# 3. Erstelle die UI (OHNE 'theme' Argument, das den Fehler verursacht hat)
interface = gr.Interface(
    fn=generate_text,
    inputs=gr.Textbox(lines=2, placeholder="Gib hier deinen Text ein...", label="Eingabe"),
    outputs=gr.Textbox(label="Antwort"),
    title=title,
    description=description
)

# 4. Starten
if __name__ == "__main__":
    interface.launch()