wehub-resource-sync
280 lines
10 KiB
Python
280 lines
10 KiB
Python
import logging
|
|
from typing import Literal, Any
|
|
import argparse
|
|
import json
|
|
import PyPDF2
|
|
import random
|
|
import os, shutil
|
|
from math import ceil
|
|
from datasets import Dataset, concatenate_datasets
|
|
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForQuestionAnswering
|
|
import torch
|
|
# Set up logging
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger("huggingface_script")
|
|
|
|
# Document type literals
|
|
DocType = Literal["api", "pdf", "json", "txt"]
|
|
|
|
# Every N chunks, save a checkpoint
|
|
N = 15
|
|
|
|
def get_args() -> argparse.Namespace:
|
|
"""
|
|
Parses and returns the command line arguments specified by the user.
|
|
"""
|
|
parser = argparse.ArgumentParser()
|
|
|
|
parser.add_argument("--datapath", type=str, default="", help="The path at which the document is located")
|
|
parser.add_argument("--output", type=str, default="./", help="The path at which to save the dataset")
|
|
parser.add_argument("--output-format", type=str, default="hf", help="Format to convert the dataset to. Defaults to hf.")
|
|
parser.add_argument("--output-type", type=str, default="jsonl", help="Type to export the dataset to. Defaults to jsonl.")
|
|
parser.add_argument("--distractors", type=int, default=3, help="The number of distractor documents to include per data point / triplet")
|
|
parser.add_argument("--p", type=float, default=1.0, help="The percentage that the oracle document is included in the context")
|
|
parser.add_argument("--questions", type=int, default=5, help="The number of data points / triplets to generate per chunk")
|
|
parser.add_argument("--chunk_size", type=int, default=512, help="The size of each chunk in number of tokens")
|
|
parser.add_argument("--doctype", type=str, default="pdf", help="The type of the document", choices=["pdf", "txt", "json", "api"])
|
|
parser.add_argument("--fast", action="store_true", help="Run the script in fast mode (no recovery implemented)")
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
def get_chunks(file_path: str, doctype: DocType = "pdf", chunk_size: int = 512) -> list[str]:
|
|
"""
|
|
Takes in a `file_path` and `doctype`, retrieves the document, breaks it down into chunks of size
|
|
`chunk_size`, and returns the chunks as a list of strings.
|
|
"""
|
|
chunks = []
|
|
|
|
logger.info(f"Retrieving chunks from {file_path} of type {doctype}")
|
|
|
|
if doctype == "api":
|
|
# Load API documentation and process it
|
|
with open(file_path) as f:
|
|
api_docs_json = json.load(f)
|
|
chunks = [str(api_doc_json) for api_doc_json in api_docs_json]
|
|
|
|
else:
|
|
if doctype == "json":
|
|
# Load JSON document
|
|
with open(file_path, 'r') as f:
|
|
data = json.load(f)
|
|
text = data["text"]
|
|
elif doctype == "pdf":
|
|
# Load PDF and extract text
|
|
text = ""
|
|
with open(file_path, 'rb') as file:
|
|
reader = PyPDF2.PdfReader(file)
|
|
num_pages = len(reader.pages)
|
|
for page_num in range(num_pages):
|
|
page = reader.pages[page_num]
|
|
text += page.extract_text()
|
|
elif doctype == "txt":
|
|
# Load plain text document
|
|
with open(file_path, 'r') as file:
|
|
text = file.read()
|
|
else:
|
|
raise TypeError("Document is not one of the accepted types: api, pdf, json, txt")
|
|
|
|
# Split the text into chunks
|
|
num_chunks = ceil(len(text) / chunk_size)
|
|
logger.info(f"Splitting text into {num_chunks} chunks.")
|
|
for i in range(0, len(text), chunk_size):
|
|
chunks.append(text[i:i + chunk_size])
|
|
|
|
return chunks
|
|
|
|
def generate_instructions_hf(chunk: str, x: int = 5, model_name: str = "t5-small") -> list[str]:
|
|
"""
|
|
Uses a Hugging Face model to generate `x` questions based on the given text chunk, utilizing the GPU if available.
|
|
"""
|
|
# Load the Hugging Face model and tokenizer for question generation
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
|
|
|
|
# Move model to GPU if available
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
model.to(device)
|
|
|
|
input_text = f"Generate questions based on the following text: {chunk}"
|
|
inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding="longest").to(device)
|
|
|
|
outputs = model.generate(
|
|
inputs.input_ids,
|
|
max_length=64,
|
|
num_beams=x, # Using beam search with `x` beams
|
|
num_return_sequences=x # Returning `x` sequences
|
|
)
|
|
|
|
questions = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
|
|
|
|
return questions
|
|
|
|
def generate_label_hf(question: str, context: str, model_name: str = "deepset/roberta-base-squad2") -> str:
|
|
"""
|
|
Uses a Hugging Face model to generate an answer to the given question based on the context, utilizing the GPU if available.
|
|
"""
|
|
# Load the Hugging Face model and tokenizer for question-answering
|
|
question_answering_pipeline = pipeline("question-answering", model=model_name, device=0 if torch.cuda.is_available() else -1)
|
|
|
|
result = question_answering_pipeline(question=question, context=context)
|
|
|
|
return result['answer']
|
|
|
|
def add_chunk_to_dataset(
|
|
chunks: list[str],
|
|
chunk: str,
|
|
doctype: DocType = "api",
|
|
x: int = 5,
|
|
num_distract: int = 3,
|
|
p: float = 0.8,
|
|
model_name_qg: str = "t5-small",
|
|
model_name_qa: str = "deepset/roberta-base-squad2"
|
|
) -> None:
|
|
"""
|
|
Given a chunk, create {Q, A, D} triplets and add them to the dataset using Hugging Face models.
|
|
"""
|
|
global ds
|
|
i = chunks.index(chunk)
|
|
|
|
# Use the Hugging Face model to generate questions
|
|
qs = generate_instructions_hf(chunk, x, model_name=model_name_qg)
|
|
for q in qs:
|
|
datapt = {
|
|
"id": None,
|
|
"type": None,
|
|
"question": None,
|
|
"context": None,
|
|
"oracle_context": None,
|
|
"cot_answer": None
|
|
}
|
|
|
|
datapt["id"] = f"seed_task_{0 if not ds else ds.num_rows}"
|
|
datapt["type"] = "api call" if doctype == "api" else "general"
|
|
datapt["question"] = q
|
|
|
|
# Create distractor documents
|
|
docs = [chunk]
|
|
indices = list(range(0, len(chunks)))
|
|
indices.remove(i)
|
|
for j in random.sample(indices, num_distract):
|
|
docs.append(chunks[j])
|
|
# Decide whether to add oracle document
|
|
oracle = random.uniform(0, 1) < p
|
|
if not oracle:
|
|
docs[0] = chunks[random.sample(indices, 1)[0]]
|
|
random.shuffle(docs)
|
|
|
|
d = {
|
|
"title": ["placeholder_title"] * (num_distract + 1),
|
|
"sentences": docs
|
|
}
|
|
datapt["context"] = d
|
|
datapt["oracle_context"] = chunk
|
|
|
|
# Add the answer generated by the Hugging Face model
|
|
datapt["cot_answer"] = generate_label_hf(q, chunk, model_name=model_name_qa)
|
|
|
|
# Construct model instruction
|
|
context = ""
|
|
for doc in docs:
|
|
context += "<DOCUMENT>" + str(doc) + "</DOCUMENT>\n"
|
|
context += q
|
|
datapt["instruction"] = context
|
|
|
|
# Add to dataset
|
|
if not ds:
|
|
# Initialize dataset
|
|
datapt["id"] = [datapt["id"]]
|
|
datapt["type"] = [datapt["type"]]
|
|
datapt["question"] = [datapt["question"]]
|
|
datapt["context"] = [datapt["context"]]
|
|
datapt["oracle_context"] = [datapt["oracle_context"]]
|
|
datapt["cot_answer"] = [datapt["cot_answer"]]
|
|
datapt["instruction"] = [datapt["instruction"]]
|
|
ds = Dataset.from_dict(datapt)
|
|
else:
|
|
ds = ds.add_item(datapt)
|
|
|
|
def save_checkpoint(state, filename):
|
|
"""
|
|
Saves the current state of processing to a file for recovery.
|
|
"""
|
|
with open(filename, 'w') as f:
|
|
f.write(str(state))
|
|
|
|
def load_checkpoint(filename):
|
|
"""
|
|
Loads the processing state from a checkpoint file.
|
|
"""
|
|
with open(filename, 'r') as f:
|
|
return int(f.read())
|
|
|
|
def main():
|
|
global ds
|
|
|
|
# Get command line arguments
|
|
args = get_args()
|
|
|
|
CHUNK_SIZE = args.chunk_size
|
|
NUM_DISTRACT_DOCS = args.distractors
|
|
|
|
# Split the document into chunks
|
|
chunks = get_chunks(args.datapath, args.doctype, CHUNK_SIZE)
|
|
|
|
ds = None
|
|
|
|
num_chunks = len(chunks)
|
|
|
|
if not args.fast:
|
|
start = 0
|
|
if os.path.exists("checkpoint.txt"):
|
|
start = int(load_checkpoint("checkpoint.txt"))
|
|
|
|
for i in range((start // N) * N, len(chunks)):
|
|
chunk = chunks[i]
|
|
save_checkpoint(i, "checkpoint.txt")
|
|
|
|
perc = ceil(i / num_chunks * 100)
|
|
logger.info(f"Adding chunk {i}/{num_chunks}")
|
|
add_chunk_to_dataset(chunks, chunk, args.doctype, args.questions, NUM_DISTRACT_DOCS)
|
|
|
|
if (i + 1) % N == 0:
|
|
ds.save_to_disk(args.output + "-checkpoints-" + str(i))
|
|
ds = None
|
|
|
|
if ds:
|
|
ds.save_to_disk(args.output + "-checkpoints-last")
|
|
|
|
ds_list = []
|
|
|
|
for filename in os.listdir(os.path.dirname(args.output)):
|
|
if "-checkpoints-" in filename:
|
|
for f in os.listdir(os.path.dirname(args.output) + "/" + filename):
|
|
if f.endswith(".arrow"):
|
|
ds_list.append(Dataset.from_file(os.path.dirname(args.output) + "/" + filename + "/" + f))
|
|
|
|
ds = concatenate_datasets(ds_list)
|
|
else:
|
|
for i, chunk in enumerate(chunks):
|
|
perc = ceil(i / num_chunks * 100)
|
|
logger.info(f"Adding chunk {i}/{num_chunks}")
|
|
add_chunk_to_dataset(chunks, chunk, args.doctype, args.questions, NUM_DISTRACT_DOCS)
|
|
|
|
# Save the final dataset
|
|
ds.save_to_disk(args.output)
|
|
|
|
# Save as .jsonl format (dummy functionality)
|
|
# Implement a conversion function if needed, this is just a placeholder
|
|
logger.info("Converting dataset to the desired format...")
|
|
|
|
if not args.fast:
|
|
os.remove("checkpoint.txt")
|
|
for filename in os.listdir(os.path.dirname(args.output)):
|
|
if "-checkpoints-" in filename:
|
|
shutil.rmtree(os.path.dirname(args.output) + "/" + filename)
|
|
|
|
if __name__ == "__main__":
|
|
logger.info("Starting the Hugging Face processing script...")
|
|
main()
|