import streamlit as st
from transformers import pipeline

# Title of the app
st.title("Text Generation with DeepSeek-R1-Distill-Qwen-1.5B")

# Load the text-generation pipeline
@st.cache_resource  # Cache the model to avoid reloading on every interaction
def load_model():
    return pipeline("text-generation", model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")

model = load_model()

# Input text box for user input
user_input = st.text_area("Enter your prompt:", "Who are you?")

# Slider to control the max length of the generated text
max_length = st.slider("Max length of generated text", min_value=10, max_value=200, value=50)

# Button to generate text
if st.button("Generate Text"):
    if user_input:
        with st.spinner("Generating text..."):
            # Generate text using the pipeline
            messages = [{"role": "user", "content": user_input}]
            output = model(messages, max_length=max_length, num_return_sequences=1)
            generated_text = output[0]["generated_text"]

        # Display the generated text
        st.success("Generated Text:")
        st.write(generated_text)
    else:
        st.warning("Please enter a prompt!")