QA generation with Mistral Instruct v1.py

# the code below was run on google colab
# it allows for Question and Answer generation from a text file using Mistral Instruct v1
# convert your document to a text file > create your prompts


!pip install transformers
!pip install pydantic
!pip install tqdm

import json
from tqdm import tqdm
from pydantic import BaseModel, Field, ValidationError
from transformers import AutoModelForCausalLM, AutoTokenizer

from huggingface_hub import login
login(token="your token") # generate this token on huggingface

from google.colab import drive

# Mount Google Drive
drive.mount('/content/drive')

# Define your prompts
prompts = {
    "xx-1": """
    xxx.....
    Here is the document:
    {{document_chunk}}
    """,
    "xx-2": """
    xxx.....
    Here is the document:
    {{document_chunk}}
    """
}

# Define the schema for QA pairs
class QAPair(BaseModel):
    question: str = Field(description="Generated question")
    answer: str = Field(description="Generated answer")

# Function to generate and validate JSON response
def generate_response(prompt, max_length=1024):
    inputs = tokenizer(prompt, return_tensors='pt')
    input_ids = inputs['input_ids']
    attention_mask = inputs['attention_mask']
    outputs = model.generate(
        input_ids,
        attention_mask=attention_mask,
        max_length=max_length,
        num_return_sequences=1,
        do_sample=True,  # Enable sampling
        pad_token_id=tokenizer.eos_token_id
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

def extract_json_from_response(response):
    start_idx = response.find("[")
    end_idx = response.rfind("]") + 1
    if start_idx == -1 or end_idx == -1:
        raise ValueError("Valid JSON array not found in the response")
    json_str = response[start_idx:end_idx]
    return json_str


# Function to generate QA pairs using the model and validate JSON
# error validation does not seem to work but also does not seem to matter - need to fix
def generate_qa_pairs(prompt_template, document_chunk, num_questions):
    chunk_text = tokenizer.decode(document_chunk, skip_special_tokens=True)
    prompt = prompt_template.replace("{{document_chunk}}", chunk_text)
    responses = []

    for _ in range(num_questions):
        response = generate_response(prompt)
        try:
            json_str = extract_json_from_response(response)
            qa_pairs = json.loads(json_str)
            for qa in qa_pairs:
                qa_model = QAPair(**qa)  # Validate against the schema
                responses.append(qa_model.dict())
        except (json.JSONDecodeError, ValidationError, ValueError) as e:
            print(f"Error parsing or validating JSON: {e}")

    return responses

# Instantiate the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained('mistralai/Mistral-7B-Instruct-v0.1')
model = AutoModelForCausalLM.from_pretrained('mistralai/Mistral-7B-Instruct-v0.1')

# Load your document from a text file
with open('yourtextfile.txt', 'r') as file:
    document = file.read()

# Parameters
chunk_size = 512  # Adjust as needed
overlap = 128  # Adjust as needed
num_questions = 5  # Adjust as needed
use_all_prompts = False  # Set to False if you want to use a single prompt
selected_prompt = "xx-2"  # Set the prompt you want to use if not using all


# Function to chunk the document based on tokens
def chunk_document(doc, chunk_size, overlap):
    tokens = tokenizer(doc, return_tensors='pt')['input_ids'][0]
    chunks = []
    for i in range(0, len(tokens), chunk_size - overlap):
        chunks.append(tokens[i:i + chunk_size])
    return chunks

# Chunk the document
document_chunks = chunk_document(document, chunk_size, overlap)

# Dictionary to store QA pairs
all_qa_pairs = {}

# Wrap the outer loop with tqdm to show progress
for chunk in tqdm(document_chunks, desc="Processing chunks"):
    if use_all_prompts:
        for prompt_name, prompt_template in prompts.items():
            qa_pairs = generate_qa_pairs(prompt_template, chunk, num_questions)
            if prompt_name not in all_qa_pairs:
                all_qa_pairs[prompt_name] = []
            all_qa_pairs[prompt_name].extend(qa_pairs)
    else:
        qa_pairs = generate_qa_pairs(prompts[selected_prompt], chunk, num_questions)
        if selected_prompt not in all_qa_pairs:
            all_qa_pairs[selected_prompt] = []
        all_qa_pairs[selected_prompt].extend(qa_pairs)

# Save the generated QA pairs to a JSON file
with open('yourfilename.json', 'w') as outfile:
    json.dump(all_qa_pairs, outfile, indent=2)

print("QA pairs have been saved to qa_pairs file")

# Define the path to save the file in Google Drive
output_path = '/content/drive/My Drive/yourfilename.json'

# Save the generated QA pairs to a JSON file in Google Drive
with open(output_path, 'w') as outfile:
    json.dump(all_qa_pairs, outfile, indent=2)

print(f"QA pairs have been saved to {output_path}")