Spaces:
Sleeping
Sleeping
import gradio as gr | |
from transformers import AutoModelForCausalLM, AutoTokenizer | |
import fitz # PyMuPDF | |
import traceback | |
# Carregar o tokenizer e o modelo | |
tokenizer = AutoTokenizer.from_pretrained("Locutusque/gpt2-xl-conversational") | |
model = AutoModelForCausalLM.from_pretrained("Locutusque/gpt2-xl-conversational") | |
# Variável global para armazenar o conteúdo do PDF | |
pdf_content = "" | |
# Função para ler o PDF | |
def read_pdf(file_path): | |
try: | |
doc = fitz.open(file_path) | |
text = "" | |
for page_num in range(len(doc)): | |
page = doc.load_page(page_num) | |
text += page.get_text() | |
return text | |
except Exception as e: | |
print(f"Erro ao ler o PDF: {e}") | |
print(traceback.format_exc()) | |
return "" | |
# Função para carregar o PDF e armazenar o conteúdo | |
def load_pdf(pdf_file): | |
global pdf_content | |
try: | |
pdf_path = pdf_file.name # Acessar o caminho do arquivo carregado | |
pdf_content = read_pdf(pdf_path) | |
if pdf_content: | |
return "PDF carregado com sucesso!" | |
else: | |
return "Falha ao carregar o PDF." | |
except Exception as e: | |
print(f"Erro ao carregar o PDF: {e}") | |
print(traceback.format_exc()) | |
return "Erro ao carregar o PDF." | |
# Função para responder perguntas com base no conteúdo do PDF | |
def answer_question(question, max_length=200, temperature=0.7, top_k=50, top_p=0.95): | |
global pdf_content | |
try: | |
if not pdf_content: | |
return "Por favor, carregue um PDF primeiro.", "" | |
# Limitar o tamanho do conteúdo do PDF para evitar exceder o número máximo de tokens | |
max_tokens = 512 # Limitar o número de tokens do conteúdo do PDF para evitar sobrecarga | |
truncated_content = pdf_content[:max_tokens] | |
prompt = f"Conteúdo do PDF: {truncated_content}\nPergunta: {question}\nResposta em português:" | |
# Verificando o tamanho do prompt | |
print(f"Prompt gerado (tamanho {len(prompt)}):\n{prompt[:500]}...") # Imprime os primeiros 500 caracteres para depuração | |
# Tokenizando o prompt | |
inputs = tokenizer(prompt, return_tensors="pt") | |
num_tokens = inputs.input_ids.shape[1] | |
print(f"Tokens de entrada: {num_tokens} tokens") # Verifica o número de tokens | |
if num_tokens > 1024: | |
return "Erro: O prompt gerado ultrapassou o limite de tokens do modelo. Tente um PDF menor.", prompt | |
# Gerar a resposta com configurações ajustadas | |
print(f"Gerando resposta com max_new_tokens=100 e temperature={temperature}") | |
outputs = model.generate( | |
inputs.input_ids, | |
max_new_tokens=100, # Gerar até 100 novos tokens (apenas a resposta) | |
temperature=temperature, # Temperatura para controlar a criatividade | |
num_return_sequences=1, | |
no_repeat_ngram_size=2, # Evita repetições excessivas | |
pad_token_id=tokenizer.eos_token_id, # Adiciona um token de padding adequado | |
) | |
# Decodificando a resposta | |
response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
print(f"Resposta gerada: {response}") | |
return response, prompt | |
except Exception as e: | |
print(f"Erro ao gerar resposta: {e}") | |
print(traceback.format_exc()) | |
return f"Erro ao gerar resposta: {e}", "" | |
# Interface do Gradio para carregar PDF e fazer perguntas | |
pdf_loader = gr.Interface( | |
fn=load_pdf, | |
inputs=gr.File(label="Carregue um PDF"), | |
outputs="text", | |
title="Carregar PDF" | |
) | |
question_answerer = gr.Interface( | |
fn=answer_question, | |
inputs=gr.Textbox(lines=2, label="Pergunta"), | |
outputs=[gr.Textbox(label="Resposta"), gr.Textbox(label="Prompt usado")], | |
title="Perguntas sobre o PDF" | |
) | |
# Combinar as interfaces em uma aplicação | |
iface = gr.TabbedInterface( | |
[pdf_loader, question_answerer], | |
["Carregar PDF", "Fazer Perguntas"] | |
) | |
if __name__ == "__main__": | |
iface.launch() |