import streamlit as st
from transformers import pipeline

# Load the pipeline for text generation
model_name = "mostafaHaydar/model_test"  # Replace with your model's name
generator = pipeline("text-generation", model=model_name)

# Set up the Streamlit app
st.title("Text Generation with LLaMA 3")

# Text input from the user
prompt = st.text_area("Enter your prompt:")

# Generate text when the user clicks the button
if st.button("Generate"):
    if prompt:
        # Generate text using the pipeline
        generated_text = generator(prompt, max_length=150, num_return_sequences=1)[0]['generated_text']

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