Update app.py

app.py

@@ -36,11 +36,19 @@ h1 {
 def chat_llama3_8b(message: str, history: list, temperature: float, max_new_tokens: int) -> str:
     # Initialize the conversation with a system prompt
     conversation = [{"role": "system", "content": "Sei un assistente specializzato nella lingua italiana."}]
-    
-    # Add historical conversation
-    for user, assistant in history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-    
+
+    flat_history = [item for sublist in history for item in sublist]
+
+    if len(flat_history) > 16:
+        flat_history = flat_history[-16:]
+
+    # Rebuild the conversation from the trimmed history
+    for i in range(0, len(flat_history), 2):
+        conversation.extend([
+            {"role": "user", "content": flat_history[i]},
+            {"role": "assistant", "content": flat_history[i + 1]}
+        ])
+
     # Add the current user message to the conversation
     conversation.append({"role": "user", "content": message})
 
@@ -59,13 +67,11 @@ def chat_llama3_8b(message: str, history: list, temperature: float, max_new_toke
         temperature=real_temperature,
         eos_token_id=tokenizer.eos_token_id
     )
+    input_length = input_ids.size(1)
+    new_tokens = generated_ids[:, input_length:]
+    decoded = tokenizer.batch_decode(new_tokens, skip_special_tokens=True)[0]
 
-    # Decode the generated tokens
-    decoded = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
-    prompt_end_index = decoded[0].find(message) + len(message)
-    final_response = decoded[0][prompt_end_index:] if prompt_end_index != -1 else decoded[0]
-    
-    final_response = final_response.strip("assistant")
+    final_response = decoded.strip("assistant")
     if final_response.startswith(':'):
         final_response = final_response.lstrip(':').strip()