Partial bot response on user's stop generation

app.py

@@ -5,15 +5,24 @@ from streamlit import session_state as ss
 import os
 import login  # Importa il file login.py che hai creato
 
+# Stile del chat input e sfondo della pagina
 st.markdown("""
     <style>
-    /* Cambia il colore del testo all'interno del chat input */
     section[data-testid="stTextInput"] input {
-        color: black !important; /* Cambia 'black' al colore che preferisci, es: #FFFFFF per bianco */
-        background-color: #F0F2F6 !important; /* Cambia il colore dello sfondo del chat input */
-        font-size: 16px; /* Cambia la dimensione del testo se necessario */
-        border-radius: 10px; /* Arrotonda i bordi del chat input */
-        padding: 10px; /* Aggiungi spazio interno per rendere l'aspetto più pulito */
+        color: black !important;
+        background-color: #F0F2F6 !important;
+        font-size: 16px;
+        border-radius: 10px;
+        padding: 10px;
+    }
+    .main {
+        background-color: #0A0A1A;
+        color: #FFFFFF;
+    }
+    .stChatMessage div[data-baseweb="block"] {
+        background-color: rgba(255, 255, 255, 0.1) !important;
+        color: #FFFFFF !important;
+        border-radius: 10px !important;
     }
     </style>
 """, unsafe_allow_html=True)
@@ -24,14 +33,13 @@ if "is_logged_in" not in st.session_state:
 
 # Mostra la pagina di login solo se l'utente non è loggato
 if not st.session_state["is_logged_in"]:
-    login.login_page()  # Mostra la pagina di login e blocca il caricamento della chat
+    login.login_page()
     st.stop()
 
 # Recupera le secrets da Hugging Face
-model_repo = st.secrets["MODEL_REPO"]  # Repository del modello di base
-hf_token = st.secrets["HF_TOKEN"]  # Token Hugging Face
+model_repo = st.secrets["MODEL_REPO"]
+hf_token = st.secrets["HF_TOKEN"]
 
-# Carica il modello di base con caching
 @st.cache_resource
 def load_model():
     tokenizer = AutoTokenizer.from_pretrained(model_repo, use_auth_token=hf_token)
@@ -45,9 +53,9 @@ def generate_llama_response_stream(user_input, tokenizer, model, max_length=512)
     input_ids = tokenizer.encode(user_input + eos_token, return_tensors="pt")
     response_text = ""
 
-    response_placeholder = st.empty()  # Placeholder per mostrare la risposta progressiva
+    response_placeholder = st.empty()
 
-    # Genera un token alla volta e aggiorna il placeholder
+    # Genera un token alla volta e aggiorna il testo in response_text
     for i in range(max_length):
         if ss.get("stop_generation", False):
             break  # Interrompe il ciclo se l'utente ha premuto "stop"
@@ -57,7 +65,10 @@ def generate_llama_response_stream(user_input, tokenizer, model, max_length=512)
         new_token = tokenizer.decode([new_token_id], skip_special_tokens=True)
 
         response_text += new_token
-        response_placeholder.markdown(f"**RacoGPT:** {response_text}")  # Aggiorna il testo progressivo
+        response_placeholder.markdown(f"RacoGPT: {response_text}", unsafe_allow_html=True)
+
+        # Salva il testo parziale in session_state per preservarlo in caso di interruzione
+        ss["response_text_partial"] = response_text
 
         # Aggiungi il nuovo token alla sequenza di input
         input_ids = torch.cat([input_ids, output[:, -1:]], dim=-1)
@@ -86,39 +97,14 @@ if 'stop_generation' not in ss:
 # Carica il modello e tokenizer
 tokenizer, model = load_model()
 
-# Aggiungi il CSS per personalizzare il colore del testo e lo sfondo del chat input
-st.markdown("""
-    <style>
-    /* Sfondo e colore del testo del chat input */
-    section[data-testid="stTextInput"] input {
-        color: black !important;
-        background-color: #F0F2F6 !important;
-        border-radius: 10px !important;
-    }
-    
-    /* Sfondo scuro per l'intera pagina */
-    .main {
-        background-color: #0A0A1A;
-        color: #FFFFFF;
-    }
-
-    /* Modifica sfondo e colore dei messaggi di chat */
-    .stChatMessage div[data-baseweb="block"] {
-        background-color: rgba(255, 255, 255, 0.1) !important;
-        color: #FFFFFF !important;
-        border-radius: 10px !important;
-    }
-    </style>
-""", unsafe_allow_html=True)
-
 # Mostra la cronologia dei messaggi con le label personalizzate
 for message in ss.msg:
     if message["role"] == "user":
         with st.chat_message("user"):
-            st.markdown(f"**Tu:** {message['content']}")
+            st.markdown(f"Tu: {message['content']}")
     elif message["role"] == "RacoGPT":
         with st.chat_message("RacoGPT"):
-            st.markdown(f"**RacoGPT:** {message['content']}")
+            st.markdown(f"RacoGPT: {message['content']}")
 
 # Contenitore per gestire la mutua esclusione tra input e pulsante di stop
 input_container = st.empty()
@@ -129,11 +115,10 @@ if not ss.is_chat_input_disabled:
         prompt = st.chat_input("Scrivi il tuo messaggio...")
 
     if prompt:
-        # Salva il messaggio dell'utente e disabilita l'input
         ss.msg.append({"role": "user", "content": prompt})
         with st.chat_message("user"):
             ss.is_chat_input_disabled = True
-            st.markdown(f"**Tu:** {prompt}")
+            st.markdown(f"Tu: {prompt}")
             st.rerun()
 else:
     # Mostra il pulsante di "Stop Generazione" al posto della barra di input
@@ -145,10 +130,16 @@ else:
     with st.spinner("RacoGPT sta generando una risposta..."):
         response = generate_llama_response_stream(ss.msg[-1]['content'], tokenizer, model)
 
-    # Mostra il messaggio finale del bot dopo che la risposta è completata o interrotta
-    ss.msg.append({"role": "RacoGPT", "content": response})
+    # Usa il testo parziale se presente
+    final_response = response or ss.get("response_text_partial", "")
+
+    # Aggiungi la risposta finale nella cronologia dei messaggi
+    ss.msg.append({"role": "RacoGPT", "content": final_response})
     with st.chat_message("RacoGPT"):
-        st.markdown(f"**RacoGPT:** {response}")
+        st.markdown(f"RacoGPT: {final_response}")
+
+    # Pulisce il testo parziale dalla sessione e riabilita l'input
+    ss.pop("response_text_partial", None)
     ss.is_chat_input_disabled = False
 
     # Rerun per aggiornare l'interfaccia