Update app.py

app.py

@@ -32,36 +32,40 @@ def chat(history, system_message, max_tokens, temperature, top_p, top_k, repetit
     history = history or []
     
     # A última mensagem do usuário
-    user_prompt = history[-1][0]
+    user_prompt = history[-1][0] if history else ""
     
-    # Prepare the messages
-    messages = [
-        {"role": "system", "content": system_message.strip()},
-        {"role": "user", "content": user_prompt}
-    ]
+    # Preparar a entrada para o modelo
+    prompt_template = f'''system
+{system_message.strip()}
+user
+{user_prompt}
+assistant
+'''
+    input_ids = tokenizer(prompt_template, return_tensors='pt').input_ids
     
-    # Apply the chat template
-    gen_input = tokenizer.apply_chat_template(messages, return_tensors="pt", add_generation_prompt=True)
-    
-    # Debug: Print the shape of gen_input
-    print("Shape of gen_input:", gen_input.shape if hasattr(gen_input, 'shape') else type(gen_input))
-    
-    # Extract input_ids based on the type of gen_input
-    input_ids = gen_input['input_ids'] if isinstance(gen_input, dict) else gen_input
-    
-    # Generate the output
-    output = model.generate(input_ids=input_ids, temperature=temperature, do_sample=True, top_p=top_p, top_k=top_k, max_length=max_tokens)
+    # Gerar a saída
+    output = model.generate(
+        input_ids=input_ids,
+        max_length=max_tokens,
+        temperature=temperature,
+        top_p=top_p,
+        top_k=top_k,
+        repetition_penalty=repetition_penalty
+    )
     
-    # Decode the output
-    decoded_output = tokenizer.decode(output[0], skip_special_tokens=True)
+    # Decodificar a saída
+    decoded_output = tokenizer.decode(output[0])
+    assistant_response = decoded_output.split('assistant')[-1].strip()  # Pegar apenas a última resposta do assistente
     
-    # Update the history
-    history[-1][1] += decoded_output.split('\n')[-1]  # Only take the assistant's last response
+    # Atualizar o histórico
+    if history:
+        history[-1][1] += assistant_response
+    else:
+        history.append(["", assistant_response])
     
     return history, history, ""
 
 
-
 start_message = ""
 
 CSS ="""