Upload app.py

response not displayed corrected

app.py

@@ -1,5 +1,5 @@
 import streamlit as st
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from threading import Thread
 
 # Model Initialization
@@ -28,38 +28,13 @@ def load_model():
 tokenizer, llama_pipeline = load_model()
 
 # Generate text
-def generate_response(prompt, chat_history, max_new_tokens):
-    history = []
-
-    # Build chat history
-    for sent, received in chat_history:
-        history.append({"role": "user", "content": sent})
-        history.append({"role": "assistant", "content": received})
-
-    history.append({"role": "user", "content": prompt})
-
-    if len(tokenizer.apply_chat_template(history)) > 512:
-        return "Chat history is too long."
-    else:
-        streamer = TextIteratorStreamer(
-            tokenizer=tokenizer,
-            skip_prompt=True,
-            skip_special_tokens=True,
-            timeout=300.0
-        )
-        thread = Thread(target=llama_pipeline, kwargs={
-            "text_inputs": history,
-            "max_new_tokens": max_new_tokens,
-            "repetition_penalty": 1.15,
-            "streamer": streamer
-        })
-        thread.start()
-
-        generated_text = ""
-        for word in streamer:
-            generated_text += word
-            response = generated_text.strip()
-            yield response
+def generate_response(prompt, max_new_tokens):
+    response = llama_pipeline(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        repetition_penalty=1.15
+    )
+    return response[0]['generated_text']
 
 # Sidebar: Configuration
 st.sidebar.header("Chatbot Configuration")
@@ -73,18 +48,13 @@ examples = [
     "የአማርኛ ግጥም ፃፍልኝ",
     "ተረት ንገረኝ\n\nጅብና አንበሳ",
     "አንድ አስቂኝ ቀልድ ንገረኝ",
-    "የተሰጠው ጽሑፍ አስተያየት ምን አይነት ነው? 'አዎንታዊ'፣ 'አሉታዊ' ወይም 'ገለልተኛ' የሚል ምላሽ ስጥ። 'አሪፍ ፊልም ነበር'",
     "የፈረንሳይ ዋና ከተማ ስም ምንድን ነው?",
     "አሁን የአሜሪካ ፕሬዚዳንት ማን ነው?",
-    "ሶስት የአፍሪካ ሀገራት ጥቀስልኝ",
-    "3 የአሜሪካ መሪዎችን ስም ጥቀስ",
-    "5 የአሜሪካ ከተማዎችን ጥቀስ",
-    "አምስት የአውሮፓ ሀገሮችን ጥቀስልኝ",
-    "በ ዓለም ላይ ያሉትን 7 አህጉራት ንገረኝ"
 ]
 
 st.subheader("Chat with the Amharic Chatbot")
-chat_history = st.session_state.get("chat_history", [])
+if "chat_history" not in st.session_state:
+    st.session_state.chat_history = []
 
 # Example selector
 example = st.selectbox("Choose an example:", ["Type your own message"] + examples)
@@ -94,21 +64,14 @@ user_input = st.text_input("Your message:", value=example if example != "Type yo
 
 if st.button("Send"):
     if user_input:
-        st.session_state.chat_history = st.session_state.get("chat_history", [])
-        st.session_state.chat_history.append((user_input, ""))
-        responses = generate_response(user_input, st.session_state.chat_history, max_tokens)
-
-        # Stream output
+        # Generate response
         with st.spinner("Generating response..."):
-            final_response = ""
-            for response in responses:
-                final_response = response
-                st.session_state.chat_history[-1] = (user_input, final_response)
-                st.rerun()
+            response = generate_response(user_input, max_tokens)
+            st.session_state.chat_history.append((user_input, response))
 
 # Display Chat History
-if "chat_history" in st.session_state:
-    for i, (user_msg, bot_response) in enumerate(st.session_state.chat_history):
-        st.write(f"**User {i+1}:** {user_msg}")
-        st.write(f"**Bot:** {bot_response}")
+st.write("### Chat History")
+for i, (user_msg, bot_response) in enumerate(st.session_state.chat_history):
+    st.write(f"**User {i+1}:** {user_msg}")
+    st.write(f"**Bot:** {bot_response}")