gerador_QA / app.py
igoracmorais's picture
Update app.py
40f834a verified
raw
history blame
2.64 kB
import PyPDF2
import gradio as gr
import json
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
# Função para extrair texto do PDF
def extract_text_from_pdf(pdf_file):
reader = PyPDF2.PdfReader(pdf_file)
text = ""
for page in reader.pages:
text += page.extract_text()
return text
# Função para gerar perguntas e respostas usando um modelo da Hugging Face
def generate_qa_pairs(text):
tokenizer = AutoTokenizer.from_pretrained("valhalla/t5-base-qg-hl")
model = AutoModelForSeq2SeqLM.from_pretrained("valhalla/t5-base-qg-hl")
inputs = tokenizer.encode("generate questions: " + text, return_tensors="pt", max_length=512, truncation=True)
outputs = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True)
questions = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
# O modelo retorna apenas as perguntas, então precisamos criar respostas fictícias para o exemplo
qas = [{"question": question, "answer": "answer", "answer_start": 0} for question in questions]
return qas
# Função para converter os pares de QA no formato SQuAD
def convert_to_squad_format(qas, context):
squad_data = []
for i, qa in enumerate(qas):
entry = {
"title": "Generated Data",
"context": context,
"question": qa['question'],
"id": str(i),
"answers": {
"answer_start": [qa['answer_start']],
"text": [qa['answer']]
}
}
squad_data.append(entry)
return squad_data
# Função para salvar os dados no formato SQuAD
def save_to_json(data, file_name):
if not file_name.endswith(".json"):
file_name += ".json"
with open(file_name, "w", encoding='utf-8') as f:
json.dump(data, f, ensure_ascii=False, indent=4)
return file_name
# Função principal para ser usada no Gradio
def process_pdf(pdf_file, file_name):
context = extract_text_from_pdf(pdf_file)
qas = generate_qa_pairs(context)
squad_data = convert_to_squad_format(qas, context)
file_path = save_to_json(squad_data, file_name)
return file_path
# Interface Gradio
with gr.Blocks() as demo:
with gr.Row():
pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
file_name = gr.Textbox(label="Output JSON File Name", value="squad_dataset")
process_button = gr.Button("Process PDF")
download_link = gr.File(label="Download JSON", interactive=False)
process_button.click(fn=process_pdf, inputs=[pdf_file, file_name], outputs=download_link)
demo.launch()