spedrox-sac's picture
Update app.py
b93b0a9 verified
raw
history blame
810 Bytes
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)