Update app.py

app.py

@@ -63,31 +63,31 @@ terminators = [
 ]
 
 
-def generate_response(input_ids, generate_kwargs):
+def generate_response(input_ids, generate_kwargs, output_queue):
     try:
-        # Generate the output using the model
         output = model.generate(**generate_kwargs)
-        return output
+        output_queue.append(output)
     except Exception as e:
         print(f"Error during generation: {e}")
+        output_queue.append(None)
 
         
 @spaces.GPU(duration=120)
 def chat_llama3_8b(message, history, temperature=0.95, max_new_tokens=512):
     # Prepare conversation context
-    conversation = [{"role": "user", "content": message}] + [{"role": "assistant", "content": reply} for reply in history]
-    input_ids = tokenizer(conversation, return_tensors="pt", padding=True, truncation=True).input_ids.to(model.device)
+    conversation = [message] + [msg for pair in history for msg in pair]
+    inputs = tokenizer(conversation, return_tensors="pt", padding=True, truncation=True).input_ids.to(model.device)
     
     generate_kwargs = {
-        "input_ids": input_ids,
-        "max_length": input_ids.shape[1] + max_new_tokens,
+        "input_ids": inputs,
+        "max_length": inputs.shape[1] + max_new_tokens,
         "temperature": temperature,
         "num_return_sequences": 1
     }
 
     # Thread for generating model response
     output_queue = []
-    response_thread = Thread(target=generate_response, args=(input_ids, generate_kwargs, output_queue))
+    response_thread = Thread(target=generate_response, args=(inputs, generate_kwargs, output_queue))
     response_thread.start()
     response_thread.join()  # Wait for the thread to complete