Update app.py

app.py

@@ -22,21 +22,23 @@ def response(user_question, table_data):
     # The query should be passed as a list
     encoding = tokenizer(table=table_data, queries=[user_question], padding=True, return_tensors="pt", truncation=True)
 
-    # Experiment with generation parameters
-    outputs = model.generate(
-        **encoding,
-        num_beams=5,            # Beam search to generate more diverse responses
-        top_k=50,               # Top-k sampling for diversity
-        top_p=0.95,             # Nucleus sampling
-        temperature=0.7,        # Temperature scaling (if supported by the model)
-        max_length=50,          # Limit the length of the generated response
-        early_stopping=True     # Stop generation when an end token is generated
+    # Instead of using generate, we pass the encoding through the model to get the logits
+    outputs = model(**encoding)
+
+    # Extract the answer coordinates
+    predicted_answer_coordinates = outputs.logits.argmax(-1)
+
+    # Decode the answer from the table using the coordinates
+    answer = tokenizer.convert_logits_to_predictions(
+        encoding.data,
+        predicted_answer_coordinates
     )
 
-    ans = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+    # Process the answer into a readable format
+    answer_text = answer[0][0][0] if len(answer[0]) > 0 else "Não foi possível encontrar uma resposta"
 
     query_result = {
-        "Resposta": ans[0]
+        "Resposta": answer_text
     }
 
     b = datetime.datetime.now()
@@ -44,7 +46,6 @@ def response(user_question, table_data):
 
     return query_result
 
-
 # Streamlit interface
 st.markdown("""
 <div style='display: flex; align-items: center;'>