Spaces:
Running
Running
| import streamlit as st | |
| import pandas as pd | |
| import torch | |
| from transformers import pipeline | |
| #from transformers import TapasTokenizer, TapexTokenizer, BartForConditionalGeneration | |
| from transformers import AutoTokenizer, AutoModelForTableQuestionAnswering | |
| import datetime | |
| #df = pd.read_excel('discrepantes.xlsx', index_col='Unnamed: 0') | |
| df = pd.read_excel('discrepantes.xlsx') | |
| df.fillna(0, inplace=True) | |
| table_data = df.astype(str) | |
| print(table_data.head()) | |
| def response(user_question, table_data): | |
| a = datetime.datetime.now() | |
| model_name = "google/tapas-base-finetuned-wtq" | |
| model = AutoModelForTableQuestionAnswering.from_pretrained(model_name) | |
| tokenizer = AutoTokenizer.from_pretrained(model_name) | |
| # 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 | |
| ) | |
| ans = tokenizer.batch_decode(outputs, skip_special_tokens=True) | |
| query_result = { | |
| "Resposta": ans[0] | |
| } | |
| b = datetime.datetime.now() | |
| print(b - a) | |
| return query_result | |
| # Streamlit interface | |
| st.markdown(""" | |
| <div style='display: flex; align-items: center;'> | |
| <div style='width: 40px; height: 40px; background-color: green; border-radius: 50%; margin-right: 5px;'></div> | |
| <div style='width: 40px; height: 40px; background-color: red; border-radius: 50%; margin-right: 5px;'></div> | |
| <div style='width: 40px; height: 40px; background-color: yellow; border-radius: 50%; margin-right: 5px;'></div> | |
| <span style='font-size: 40px; font-weight: bold;'>Chatbot do Tesouro RS</span> | |
| </div> | |
| """, unsafe_allow_html=True) | |
| # Chat history | |
| if 'history' not in st.session_state: | |
| st.session_state['history'] = [] | |
| # Input box for user question | |
| user_question = st.text_input("Escreva sua questão aqui:", "") | |
| if user_question: | |
| # Add person emoji when typing question | |
| st.session_state['history'].append(('👤', user_question)) | |
| st.markdown(f"**👤 {user_question}**") | |
| # Generate the response | |
| bot_response = response(user_question, table_data) | |
| # Add robot emoji when generating response and align to the right | |
| st.session_state['history'].append(('🤖', bot_response)) | |
| st.markdown(f"<div style='text-align: right'>**🤖 {bot_response}**</div>", unsafe_allow_html=True) | |
| # Clear history button | |
| if st.button("Limpar"): | |
| st.session_state['history'] = [] | |
| # Display chat history | |
| for sender, message in st.session_state['history']: | |
| if sender == '👤': | |
| st.markdown(f"**👤 {message}**") | |
| elif sender == '🤖': | |
| st.markdown(f"<div style='text-align: right'>**🤖 {message}**</div>", unsafe_allow_html=True) | |