fix app.py

app.py

@@ -230,62 +230,116 @@ def combine_history(prompt):
     return total_prompt
 
 
-def main():
-    st.title('internlm2_5-7b-chat-assistant')
-
-    # torch.cuda.empty_cache()
-    print('load model begin.')
-    model, tokenizer = load_model()
-    print('load model end.')
-
-    generation_config = prepare_generation_config()
-
-    # Initialize chat history
-    if 'messages' not in st.session_state:
-        st.session_state.messages = []
-
-    # Display chat messages from history on app rerun
-    for message in st.session_state.messages:
-        with st.chat_message(message['role'], avatar=message.get('avatar')):
-            st.markdown(message['content'])
-
-    # Accept user input
-    if prompt := st.chat_input('What is up?'):
-        # Display user message in chat message container
-
-        with st.chat_message('user', avatar='user'):
-
-            st.markdown(prompt)
-        real_prompt = combine_history(prompt)
-        # Add user message to chat history
-        st.session_state.messages.append({
-            'role': 'user',
-            'content': prompt,
-            'avatar': 'user'
-        })
-
-        with st.chat_message('robot', avatar='assistant'):
-
-            message_placeholder = st.empty()
-            for cur_response in generate_interactive(
-                    model=model,
-                    tokenizer=tokenizer,
-                    prompt=real_prompt,
-                    additional_eos_token_id=92542,
-                    device='cuda:0',
-                    **asdict(generation_config),
-            ):
-                # Display robot response in chat message container
-                message_placeholder.markdown(cur_response + '▌')
-            message_placeholder.markdown(cur_response)
-        # Add robot response to chat history
-        st.session_state.messages.append({
-            'role': 'robot',
-            'content': cur_response,  # pylint: disable=undefined-loop-variable
-            'avatar': 'assistant',
-        })
-        torch.cuda.empty_cache()
-
-
-if __name__ == '__main__':
-    main()
+# def main():
+#     st.title('internlm2_5-7b-chat-assistant')
+
+#     # torch.cuda.empty_cache()
+#     print('load model begin.')
+#     model, tokenizer = load_model()
+#     print('load model end.')
+
+#     generation_config = prepare_generation_config()
+
+#     # Initialize chat history
+#     if 'messages' not in st.session_state:
+#         st.session_state.messages = []
+
+#     # Display chat messages from history on app rerun
+#     for message in st.session_state.messages:
+#         with st.chat_message(message['role'], avatar=message.get('avatar')):
+#             st.markdown(message['content'])
+
+#     # Accept user input
+#     if prompt := st.chat_input('What is up?'):
+#         # Display user message in chat message container
+
+#         with st.chat_message('user', avatar='user'):
+
+#             st.markdown(prompt)
+#         real_prompt = combine_history(prompt)
+#         # Add user message to chat history
+#         st.session_state.messages.append({
+#             'role': 'user',
+#             'content': prompt,
+#             'avatar': 'user'
+#         })
+
+#         with st.chat_message('robot', avatar='assistant'):
+
+#             message_placeholder = st.empty()
+#             for cur_response in generate_interactive(
+#                     model=model,
+#                     tokenizer=tokenizer,
+#                     prompt=real_prompt,
+#                     additional_eos_token_id=92542,
+#                     device='cuda:0',
+#                     **asdict(generation_config),
+#             ):
+#                 # Display robot response in chat message container
+#                 message_placeholder.markdown(cur_response + '▌')
+#             message_placeholder.markdown(cur_response)
+#         # Add robot response to chat history
+#         st.session_state.messages.append({
+#             'role': 'robot',
+#             'content': cur_response,  # pylint: disable=undefined-loop-variable
+#             'avatar': 'assistant',
+#         })
+#         torch.cuda.empty_cache()
+
+
+# if __name__ == '__main__':
+#     main()
+st.title('internlm2_5-7b-chat-assistant')
+
+# torch.cuda.empty_cache()
+print('load model begin.')
+model, tokenizer = load_model()
+print('load model end.')
+
+generation_config = prepare_generation_config()
+
+# Initialize chat history
+if 'messages' not in st.session_state:
+    st.session_state.messages = []
+
+# Display chat messages from history on app rerun
+for message in st.session_state.messages:
+    with st.chat_message(message['role'], avatar=message.get('avatar')):
+        st.markdown(message['content'])
+
+# Accept user input
+if prompt := st.chat_input('What is up?'):
+    # Display user message in chat message container
+
+    with st.chat_message('user', avatar='user'):
+
+        st.markdown(prompt)
+    real_prompt = combine_history(prompt)
+    # Add user message to chat history
+    st.session_state.messages.append({
+        'role': 'user',
+        'content': prompt,
+        'avatar': 'user'
+    })
+
+    with st.chat_message('robot', avatar='assistant'):
+
+        message_placeholder = st.empty()
+        for cur_response in generate_interactive(
+                model=model,
+                tokenizer=tokenizer,
+                prompt=real_prompt,
+                additional_eos_token_id=92542,
+                device='cuda:0',
+                **asdict(generation_config),
+        ):
+            # Display robot response in chat message container
+            message_placeholder.markdown(cur_response + '▌')
+        message_placeholder.markdown(cur_response)
+    # Add robot response to chat history
+    st.session_state.messages.append({
+        'role': 'robot',
+        'content': cur_response,  # pylint: disable=undefined-loop-variable
+        'avatar': 'assistant',
+    })
+    torch.cuda.empty_cache()