Update app.py

app.py

@@ -3,8 +3,9 @@ import pandas as pd
 import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
-model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
+tokenizer = AutoTokenizer.from_pretrained("distilgpt2")
+model = AutoModelForCausalLM.from_pretrained("distilgpt2", torch_dtype=torch.float16)
+model = model.to('cuda') if torch.cuda.is_available() else model.to('cpu')
 
 # Set the padding token to the end-of-sequence token
 if tokenizer.pad_token is None:
@@ -15,23 +16,22 @@ df = pd.read_csv('anomalies.csv')
 # Função para gerar resposta
 def response(question):
     prompt = f"Considerando os dados: {df.to_string(index=False)}, onde 'ds' está em formato DateTime, 'real' é o valor da despesa e 'group' é o grupo da despesa. Pergunta: {question}"
-    inputs = tokenizer(prompt, return_tensors='pt', padding='max_length', truncation=True, max_length=512)
+    inputs = tokenizer(prompt, return_tensors='pt', padding='max_length', truncation=True, max_length=256)
     attention_mask = inputs['attention_mask']
     input_ids = inputs['input_ids']
     
     generated_ids = model.generate(
         input_ids,
         attention_mask=attention_mask,
-        max_length=len(input_ids[0]) + 100,  # Aumentar o limite de geração
-        temperature=0.65,  # Ajustar a criatividade
-        top_p=0.9,  # Usar nucleus sampling
-        no_repeat_ngram_size=2  # Evitar repetições desnecessárias
+        max_length=len(input_ids[0]) + 50,  # Reduce max_length to speed up response
+        temperature=0.7,
+        top_p=0.9,
+        no_repeat_ngram_size=2,
+        num_beams=3,  # Adding beams for more reliable generation
     )
     
     generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
-    # Processando para extrair apenas a resposta após "Resposta:"
-    response_part = generated_text.split("Resposta:")[1] if "Resposta:" in generated_text else "Resposta não encontrada."
-    final_response = response_part.split(".")[0] + "."  # Assumindo que a resposta termina na primeira sentença.
+    final_response = generated_text.split("Resposta:")[-1].split(".")[0] + "."
     
     return final_response