Update app.py

app.py

@@ -51,7 +51,7 @@ def reset_conversation():
     '''
     st.session_state.conversation = []
     st.session_state.messages = []
-    save_conversation_history([])  # Clear the conversation history file
+    save_conversation_history([])
     return None
 
 def load_conversation_history():
@@ -68,14 +68,6 @@ def save_conversation_history(conversation_history):
     with open(history_file, "wb") as f:
         pickle.dump(conversation_history, f)
 
-# Ensure session state is initialized
-if "conversation" not in st.session_state:
-    st.session_state.conversation = []
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-if "prev_option" not in st.session_state:
-    st.session_state.prev_option = None
-
 models = [key for key in model_links.keys()]
 selected_model = st.sidebar.selectbox("Select Model", models)
 temp_values = st.sidebar.slider('Select a temperature value', 0.0, 1.0, (0.5))
@@ -87,6 +79,9 @@ st.sidebar.image(model_info[selected_model]['logo'])
 
 st.sidebar.markdown("\*Generating the code might go slow if you are using low power resources \*")
 
+if "prev_option" not in st.session_state:
+    st.session_state.prev_option = selected_model
+
 if st.session_state.prev_option != selected_model:
     st.session_state.messages = []
     st.session_state.prev_option = selected_model
@@ -95,8 +90,7 @@ repo_id = model_links[selected_model]
 st.subheader(f'{selected_model}')
 
 # Load the conversation history from the file
-if not st.session_state.messages:
-    st.session_state.messages = load_conversation_history()
+st.session_state.messages = load_conversation_history()
 
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
@@ -112,13 +106,20 @@ if prompt := st.chat_input(f"Hi I'm {selected_model}, How can I help you today?"
     
     formated_text = format_promt(prompt, conversation_history, custom_instruction)
     
+    max_tokens = {
+        "LegacyLift🚀": 32000,
+        "ModernMigrate⭐": 8192,
+        "RetroRecode🔄": 4096
+    }
+    max_new_tokens = max_tokens[selected_model] - len(formated_text.split())
+    
     with st.chat_message("assistant"):
         client = InferenceClient(
             model=model_links[selected_model], )
         output = client.text_generation(
             formated_text,
             temperature=temp_values,  # 0.5
-            max_new_tokens=3000,
+            max_new_tokens=max_new_tokens,
             stream=True
         )
         response = st.write_stream(output)