Update app.py

app.py

@@ -1,21 +1,35 @@
 import streamlit as st
-from transformers import pipeline
+from transformers import pipeline, AutoTokenizer
 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)
+
+# Initialize the tokenizer and text generation pipeline
+model_name = "Qwen/Qwen2.5-0.5B-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+pipe = pipeline("text-generation", model=model_name, device=-1)
 parser = StrOutputParser()
+
 # Streamlit app
-st.title("Text Generation with Qwen Model")
+st.title("Optimized Batch Text Generation with Qwen Model")
 
 # Text input from the user
-user_input = st.text_input("Enter your message:", "Who are you?")
+user_input = st.text_area("Enter your messages (one per line):", "Who are you?\nWhat is your purpose?")
 
 # 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)
+    # Split input into multiple messages
+    user_messages = user_input.splitlines()
+    messages = [message.strip() for message in user_messages if message.strip()]
+
+    # Tokenize messages
+    tokenized_inputs = [tokenizer(message, return_tensors='pt') for message in messages]
+
+    # Process tokenized messages in a batch
+    outputs = [pipe(input_ids=tokenized_input['input_ids'], max_new_tokens=50)[0] for tokenized_input in tokenized_inputs]
+
+    # Display the generated text for each input message
+    st.write("Generated Responses:")
+    for i, output in enumerate(outputs):
+        generated_text = output['generated_text']
+        result = parser.invoke(generated_text)
+        st.write(f"Input {i+1}: {messages[i]}")
+        st.write(f"Response {i+1}: {result}\n")