import gradio as gr
from transformers import pipeline, set_seed

# Funkcja do ładowania modelu
def load_model():
    try:
        model = pipeline("text-generation", model="speakleash/Bielik-7B-v0.1")
        return model
    except Exception as e:
        print(f"Error loading model: {e}")
        return None

model = load_model()

def generate_description(product_name, product_features):
    if model is None:
        return "Model not loaded correctly."

    prompt = (
        f"Opisz produkt o nazwie {product_name}, który posiada następujące cechy: {product_features}. "
        "Podaj jego główne zalety, zastosowania i cechy wyróżniające. Wyjaśnij, dlaczego warto go kupić."
    )
    try:
        result = model(prompt, max_length=200, temperature=0.7, top_p=0.9)
        return result[0]['generated_text']
    except Exception as e:
        return f"Error generating description: {e}"

# Konfiguracja interfejsu Gradio
iface = gr.Interface(
    fn=generate_description,
    inputs=["text", "text"],
    outputs="text",
    title="Bielik-7B Product Description Generator",
    description="Wprowadź nazwę produktu i jego cechy, aby wygenerować opis."
)

if __name__ == "__main__":
    iface.launch()