Update app.py

app.py

@@ -47,54 +47,46 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # Load the model and tokenizer
 @st.cache_resource
 def load_model_and_tokenizer():
-    model_name = "microsoft/DialoGPT-medium"  # You can replace with any Hugging Face conversational model
+    model_name = "microsoft/DialoGPT-medium"  # Replace with your chosen model
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(model_name)
     return tokenizer, model
 
 tokenizer, model = load_model_and_tokenizer()
 
-# Streamlit App Title
-st.title("General Chatbot")
-st.markdown("This chatbot is powered by an open-source model from Hugging Face. Feel free to ask me anything!")
+# Streamlit App
+st.title("General Chatbot with Adjustable Answer Length")
+st.write("A chatbot powered by an open-source model from Hugging Face.")
 
-# Initialize the session state for conversation history
-if "chat_history" not in st.session_state:
-    st.session_state["chat_history"] = ""
+# Initialize the conversation
+if "conversation_history" not in st.session_state:
+    st.session_state["conversation_history"] = []
 
-# User Input Section
-user_input = st.text_input("You:", placeholder="Type your message here...", key="user_input")
+# Input for user query
+user_input = st.text_input("You:", placeholder="Ask me anything...", key="user_input")
 
-if st.button("Send") and user_input:
-    # Add user input to the conversation history
-    st.session_state["chat_history"] += f"User: {user_input}\n"
-
-    # Tokenize the input with conversation history
-    input_ids = tokenizer.encode(st.session_state["chat_history"], return_tensors="pt")
+# Slider for setting max response length
+max_length = st.slider("Set the maximum response length:", min_value=50, max_value=500, step=50, value=150)
 
-    # Generate a response
+if st.button("Send") and user_input:
+    # Append user input to history
+    st.session_state["conversation_history"].append({"role": "user", "content": user_input})
+    
+    # Tokenize and generate response
+    input_ids = tokenizer.encode(user_input + tokenizer.eos_token, return_tensors="pt")
     chat_history_ids = model.generate(
-        input_ids,
-        max_length=1500,  # Allow long responses
-        min_length=200,   # Ensure responses are not too short
-        temperature=1.0,  # Adjust for creativity
-        top_p=0.9,        # Nucleus sampling for focused responses
-        repetition_penalty=1.2,  # Penalize repeated phrases
+        input_ids, 
+        max_length=max_length,  # Use the user-specified max length
         pad_token_id=tokenizer.eos_token_id
     )
-    
-    # Decode the model's response
     response = tokenizer.decode(chat_history_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)
 
-    # Add the response to the conversation history
-    st.session_state["chat_history"] += f"Bot: {response}\n"
-
-    # Display the conversation
-    st.markdown(f"**You:** {user_input}")
-    st.markdown(f"**Bot:** {response}")
-
-# Display Full Conversation History
-st.divider()
-st.subheader("Conversation History:")
-st.text(st.session_state["chat_history"])
+    # Append model response to history
+    st.session_state["conversation_history"].append({"role": "assistant", "content": response})
 
+# Display the conversation
+for message in st.session_state["conversation_history"]:
+    if message["role"] == "user":
+        st.write(f"**You:** {message['content']}")
+    else:
+        st.write(f"**Bot:** {message['content']}")