Update app.py

app.py

@@ -3,6 +3,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import pandas as pd
 import gradio as gr
+from typing import List, Dict, Tuple, Optional
 
 # Carregando o tokenizador e o modelo
 tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-medium")
@@ -16,64 +17,67 @@ dados_pedidos = {
 df_status_pedidos = pd.DataFrame(dados_pedidos)
 
 # Função para verificar o status do pedido
-def verificar_status_pedido(numero_pedido):
-  try:
-    status = df_status_pedidos[df_status_pedidos['numero_pedido'] == numero_pedido]['status'].iloc[0]
-
-    return f'The status of your order {numero_pedido} is: {status}'
-
-  except:
-    return 'Order number not found. Please check and try again'
+def verificar_status_pedido(numero_pedido: str) -> str:
+    try:
+        status = df_status_pedidos[df_status_pedidos['numero_pedido'] == numero_pedido]['status'].iloc[0]
+        return f'The status of your order {numero_pedido} is: {status}'
+    except IndexError:
+        return 'Order number not found. Please check and try again'
 
 # Lista de palavras-chave para o status
 palavras_chave_status = ['order', 'order status', 'status of my order', 'check my order', 'track my order', 'order update']
 
-def responder(input_usuario, ids_historico_chat):
-  if any(keyword in input_usuario.lower() for keyword in palavras_chave_status):
-    return 'Could you please enter you order number?', ids_historico_chat
-  else:
-    novo_usuario_input_ids = tokenizer.encode(input_usuario + tokenizer.eos_token, return_tensors='pt')
+def responder(input_usuario: str, ids_historico_chat: Optional[torch.Tensor]) -> Tuple[str, torch.Tensor]:
+    if any(keyword in input_usuario.lower() for keyword in palavras_chave_status):
+        return 'Could you please enter your order number?', ids_historico_chat
 
+    novo_usuario_input_ids = tokenizer.encode(input_usuario + tokenizer.eos_token, return_tensors='pt')
+    
     if ids_historico_chat is not None:
-      bot_input_ids = torch.cat([ids_historico_chat, novo_usuario_input_ids ], dim=-1)
-
+        bot_input_ids = torch.cat([ids_historico_chat, novo_usuario_input_ids], dim=-1)
     else:
-      bot_input_ids = novo_usuario_input_ids
+        bot_input_ids = novo_usuario_input_ids
 
     ids_historico_chat = model.generate(
         bot_input_ids,
         max_length=1000,
         pad_token_id=tokenizer.eos_token_id
     )
-    resposta = tokenizer.decode(ids_historico_chat[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
 
-  return resposta, ids_historico_chat
+    resposta = tokenizer.decode(ids_historico_chat[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)
+    return resposta, ids_historico_chat
 
 # Criando interface
 with gr.Blocks() as app:
-  chatbot = gr.Chatbot()
-  msg = gr.Textbox(placeholder='Type your message here...')
-
-  estado = gr.State(None)
-  aguardando_numero_pedido = gr.State(False)
-
-  def processar_entrada(input_usuario, historico, ids_historico_chat, aguardando_numero_pedido):
-    if aguardando_numero_pedido:
-      resposta = verificar_status_pedido(input_usuario)
-      aguardando_numero_pedido = False
-    else:
-      resposta, ids_historico_chat = responder(input_usuario, ids_historico_chat)
-      if resposta == 'Could you please enter you order number?':
-        aguardando_numero_pedido = True
-
-    historico.append((input_usuario, resposta))
-    return historico, ids_historico_chat, aguardando_numero_pedido, ""
-
+    chatbot = gr.Chatbot(type='messages')
+    msg = gr.Textbox(placeholder='Type your message here...')
+    estado = gr.State(None)
+    aguardando_numero_pedido = gr.State(False)
+
+    def processar_entrada(input_usuario: str, historico: List[Dict[str, str]], 
+                          ids_historico_chat: Optional[torch.Tensor], 
+                          aguardando_numero_pedido: bool) -> Tuple[List[Dict[str, str]], Optional[torch.Tensor], bool, str]:
+        if aguardando_numero_pedido:
+            resposta = verificar_status_pedido(input_usuario)
+            aguardando_numero_pedido = False
+        else:
+            resposta, ids_historico_chat = responder(input_usuario, ids_historico_chat)
+        
+        if resposta == 'Could you please enter your order number?':
+            aguardando_numero_pedido = True
+        
+        historico.extend([
+            {"role": "user", "content": input_usuario},
+            {"role": "assistant", "content": resposta}
+        ])
+        
+        return historico, ids_historico_chat, aguardando_numero_pedido, ""
+
+    msg.submit(
+        processar_entrada,
+        [msg, chatbot, estado, aguardando_numero_pedido],
+        [chatbot, estado, aguardando_numero_pedido, msg]
+    )
 
-  msg.submit(
-      processar_entrada,
-      [msg, chatbot, estado, aguardando_numero_pedido],
-      [chatbot, estado, aguardando_numero_pedido, msg]
-  )
-  if "__name"=="__main":  
-      app.launch(share=True)
+if __name__ == "__main__":
+    app.launch(share=True)