Update app.py

app.py

@@ -9,10 +9,13 @@ import streamlit as st
 REPO_NAME = 'schuler/experimental-JP47D20'
 
 # Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained(REPO_NAME, trust_remote_code=True)
-generator_conf = GenerationConfig.from_pretrained(REPO_NAME)
-model = AutoModelForCausalLM.from_pretrained(REPO_NAME, trust_remote_code=True)
-generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(REPO_NAME, trust_remote_code=True)
+    generator_conf = GenerationConfig.from_pretrained(REPO_NAME)
+    model = AutoModelForCausalLM.from_pretrained(REPO_NAME, trust_remote_code=True)
+    generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
+except Exception as e:
+    st.error(f"Failed to load model: {str(e)}")
 
 # Configure the Streamlit app
 st.set_page_config(page_title="Experimental KPhi3 Model - Currently in Training", page_icon="🤗")
@@ -89,19 +92,17 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256, co
 
     Returns:
         tuple: A tuple containing the generated response and the updated chat history.
-    """
-    # Build the conversation prompt
-    prompt = "" 
-    # f"{system_message}\nCurrent Conversation:\n"
-    for message in chat_history:
-        role = "<|assistant|>" if message['role'] == 'assistant' else "<|user|>"
-        prompt += f"\n{role}\n{message['content']}\n"
-    
-
+    """    
     if continue_last:
         # We want to continue the last assistant response
-        prompt = prompt
+        prompt = st.session_state.last_response
     else:
+        # Build the conversation prompt
+        prompt = "" 
+        # f"{system_message}\nCurrent Conversation:\n"
+        for message in chat_history:
+            role = "<|assistant|>" if message['role'] == 'assistant' else "<|user|>"
+            prompt += f"\n{role}\n{message['content']}\n<|user|>\n"        
         prompt += f"\n<|user|>\n{user_text}\n<|assistant|>\n"
 
     # Generate the response
@@ -116,6 +117,8 @@ def get_response(system_message, chat_history, user_text, max_new_tokens=256, co
 
     generated_text = response_output[0]['generated_text']
 
+    st.session_state.last_response = generated_text
+
     # Extract the assistant's response
     assistant_response = generated_text[len(prompt):].strip()