import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Titel der Anwendung
st.title("🚀 Marketing Text Generator mit Falcon")
st.markdown("*Generiere kreative Marketing-Texte mit dem Falcon-7B Modell*")

# Modell und Tokenizer laden
# Wir nutzen @st.cache_resource, um zu vermeiden, dass das Modell bei jeder Interaktion neu geladen wird
@st.cache_resource
def load_falcon_model():
    """
    Lädt das Falcon Modell und den zugehörigen Tokenizer.
    Falcon ist bekannt für seine kreativen Textgenerierungsfähigkeiten.
    """
    try:
        # Modellname für Falcon
        model_name = "tiiuae/falcon-7b"
        
        # Tokenizer laden - wandelt Text in Zahlen um, die das Modell verarbeiten kann
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        
        # Modell laden mit optimierten Einstellungen für Streamlit
        model = AutoModelForCausalLM.from_pretrained(
            model_name,
            trust_remote_code=True,
            device_map="auto",     # Automatische Gerätezuweisung (CPU/GPU)
            load_in_8bit=True      # Speicheroptimierung durch 8-bit Quantisierung
        )
        
        return model, tokenizer
    
    except Exception as e:
        st.error(f"Fehler beim Laden des Modells: {str(e)}")
        return None, None

# Hauptbereich für die Eingabe
with st.form("marketing_form"):
    # Eingabefelder für die Produktinformationen
    product_name = st.text_input(
        "Produktname",
        help="Geben Sie den Namen Ihres Produkts ein"
    )
    
    key_features = st.text_area(
        "Produktmerkmale",
        help="Beschreiben Sie die wichtigsten Eigenschaften (durch Kommas getrennt)"
    )
    
    # Button zum Generieren des Textes
    submit = st.form_submit_button("Marketing-Text generieren")

# Text generieren, wenn der Button geklickt wurde
if submit and product_name and key_features:
    # Lade das Modell
    with st.spinner("Lade Falcon Modell..."):
        model, tokenizer = load_falcon_model()
    
    if model and tokenizer:
        # Marketing-spezifischer Prompt erstellen
        prompt = f"""
        Erstelle einen kurzen, überzeugenden Marketing-Text für folgendes Produkt:
        Produkt: {product_name}
        Merkmale: {key_features}
        
        Der Text sollte professionell und ansprechend sein.
        """
        
        # Text generieren
        with st.spinner("Generiere Marketing-Text..."):
            # Text in Token umwandeln
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
            
            # Generierung mit optimierten Parametern
            with torch.no_grad():  # Deaktiviert Gradient-Berechnung für Speichereffizienz
                outputs = model.generate(
                    **inputs,
                    max_length=200,         # Maximale Textlänge
                    temperature=0.7,        # Kreativität (0.0 bis 1.0)
                    top_p=0.9,             # Sampling-Parameter für Textvielfalt
                    num_return_sequences=1  # Anzahl der generierten Varianten
                )
            
            # Generierten Text dekodieren und anzeigen
            generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # Ergebnis anzeigen
            st.success("Text wurde erfolgreich generiert!")
            st.markdown("### Ihr Marketing-Text:")
            st.markdown(generated_text)
            
            # Download-Option anbieten
            st.download_button(
                "Text als Datei speichern",
                generated_text,
                file_name="marketing_text.txt"
            )
elif submit:
    st.warning("Bitte füllen Sie alle Felder aus.")

# Hilfreiche Informationen am Ende
st.markdown("---")
st.markdown("""
**Hinweise zur Nutzung:**
- Die erste Generierung kann etwas länger dauern, da das Modell geladen werden muss
- Experimentieren Sie mit verschiedenen Beschreibungen für beste Ergebnisse
""")