import streamlit as st
from transformers import pipeline
from langchain_core.output_parsers import StrOutputParser
# Initialize the text generation pipeline
pipe = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct", device=-1)
parser = StrOutputParser()
# Streamlit app
st.title("Text Generation with Qwen Model")

# Text input from the user
user_input = st.text_input("Enter your message:", "Who are you?")

# Generate text when the button is clicked
if st.button("Generate"):
    messages = [{"role": "user", "content": user_input}]
    output = pipe(messages, max_new_tokens=50)  # Adjust max_new_tokens as needed
    generated_text = output[0]['generated_text']
    result = parser.invoke(generated_text)
    # Display the generated text
    st.write("Generated Response:")
    st.write(result)