thinhlpg
2 months ago
parent
b22b02ea1d
commit
04d56325bb
@ -0,0 +1,279 @@
|
|||||||
|
"""
|
||||||
|
This script performs two main tasks:
|
||||||
|
1. It loads a markdown document, splits it into chunks, generates embeddings,
|
||||||
|
and builds a FAISS index (which is saved locally).
|
||||||
|
2. It generates QA pairs from the document using llama.
|
||||||
|
For each chunk (using a sliding window for context), it generates multiple question-answer pairs
|
||||||
|
with different difficulties. The generation is performed in batch with one retry for failed prompts.
|
||||||
|
Successfully generated QA pairs are saved to "saved_data/questions.json".
|
||||||
|
|
||||||
|
Requirements:
|
||||||
|
pip install langchain faiss-cpu unsloth vllm
|
||||||
|
"""
|
||||||
|
|
||||||
|
import json
|
||||||
|
import os
|
||||||
|
import re
|
||||||
|
from typing import Dict, List, Optional, Tuple
|
||||||
|
|
||||||
|
import pandas as pd
|
||||||
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||||
|
from loguru import logger
|
||||||
|
|
||||||
|
# Configure logger
|
||||||
|
logger.add(
|
||||||
|
"logs/generate_data_{time}.log",
|
||||||
|
rotation="500 MB",
|
||||||
|
retention="10 days",
|
||||||
|
level="INFO",
|
||||||
|
)
|
||||||
|
|
||||||
|
# ========= Part 1: Document Processing and Embedding Generation =========
|
||||||
|
# Load and split the markdown document using LangChain
|
||||||
|
from langchain_community.document_loaders import UnstructuredMarkdownLoader
|
||||||
|
from langchain_community.vectorstores import FAISS
|
||||||
|
|
||||||
|
from embeddings import CustomHuggingFaceEmbeddings
|
||||||
|
|
||||||
|
# Load your markdown file (adjust the path as needed)
|
||||||
|
loader = UnstructuredMarkdownLoader("./data/mission_report.md")
|
||||||
|
docs = loader.load()
|
||||||
|
|
||||||
|
# Split the document into smaller chunks (each 1000 characters, no overlap)
|
||||||
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
|
||||||
|
chunks = text_splitter.split_documents(docs)
|
||||||
|
|
||||||
|
# Create output directory
|
||||||
|
os.makedirs("saved_data", exist_ok=True)
|
||||||
|
|
||||||
|
# Save chunks to CSV for easy inspection
|
||||||
|
chunks_df = pd.DataFrame(
|
||||||
|
{
|
||||||
|
"chunk_id": range(1, len(chunks) + 1),
|
||||||
|
"content": [chunk.page_content for chunk in chunks],
|
||||||
|
"metadata": [chunk.metadata for chunk in chunks],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
chunks_df.to_csv("saved_data/chunks.csv", index=False)
|
||||||
|
print(f"Saved {len(chunks)} chunks to saved_data/chunks.csv")
|
||||||
|
|
||||||
|
embeddings = CustomHuggingFaceEmbeddings()
|
||||||
|
|
||||||
|
# Create a FAISS vector store from the document chunks and save it locally
|
||||||
|
vectorstore = FAISS.from_documents(chunks, embeddings)
|
||||||
|
vectorstore.save_local("faiss_index")
|
||||||
|
print("Saved FAISS index to 'faiss_index'")
|
||||||
|
|
||||||
|
# TODO: add the paraphrased chunks to the vector store
|
||||||
|
|
||||||
|
# ========= Part 2: QA Generation using Llama Backend =========
|
||||||
|
|
||||||
|
# Setup Llama backend via unsloth and vLLM
|
||||||
|
from unsloth import FastLanguageModel
|
||||||
|
from vllm import SamplingParams
|
||||||
|
|
||||||
|
import rl_helpers # Ensure you have this or remove if not used
|
||||||
|
|
||||||
|
# Load the Llama model (adjust parameters as needed)
|
||||||
|
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||||
|
model_name="meta-llama/meta-Llama-3.1-8B-Instruct",
|
||||||
|
max_seq_length=4096,
|
||||||
|
load_in_4bit=True, # Use 4-bit quantization if desired
|
||||||
|
fast_inference=True, # Enable fast inference
|
||||||
|
gpu_memory_utilization=0.6, # Adjust based on your GPU memory
|
||||||
|
)
|
||||||
|
|
||||||
|
# Define sampling parameters for generation
|
||||||
|
sampling_params = SamplingParams(
|
||||||
|
temperature=0.3,
|
||||||
|
top_p=0.95,
|
||||||
|
max_tokens=4096,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def batch_generate(prompts: List[str]) -> List[str]:
|
||||||
|
"""
|
||||||
|
Given a list of prompt strings, returns a list of generated outputs.
|
||||||
|
"""
|
||||||
|
|
||||||
|
def format_input(text: str) -> str:
|
||||||
|
return tokenizer.apply_chat_template(
|
||||||
|
[{"role": "user", "content": text}],
|
||||||
|
tokenize=False,
|
||||||
|
add_generation_prompt=True,
|
||||||
|
)
|
||||||
|
|
||||||
|
formatted = [format_input(p) for p in prompts]
|
||||||
|
outputs = model.fast_generate(formatted, sampling_params=sampling_params)
|
||||||
|
return [output.outputs[0].text for output in outputs]
|
||||||
|
|
||||||
|
|
||||||
|
def parse_qa_block(block: str) -> Optional[Tuple[str, str, str]]:
|
||||||
|
"""
|
||||||
|
Parses a QA block that should contain exactly three non-empty lines:
|
||||||
|
- A line starting with "Question:"
|
||||||
|
- A line starting with "Answer:"
|
||||||
|
- A line starting with "Difficulty:"
|
||||||
|
|
||||||
|
If the markers are not present but the block contains exactly three lines,
|
||||||
|
those are used in order.
|
||||||
|
|
||||||
|
Returns a tuple (question, answer, difficulty) or None if parsing fails.
|
||||||
|
"""
|
||||||
|
lines = [line.strip() for line in block.splitlines() if line.strip()]
|
||||||
|
if not lines:
|
||||||
|
return None
|
||||||
|
|
||||||
|
question, answer, difficulty = None, None, None
|
||||||
|
for line in lines:
|
||||||
|
lower = line.lower()
|
||||||
|
if question is None and lower.startswith("question:"):
|
||||||
|
question = line[len("question:") :].strip()
|
||||||
|
elif answer is None and lower.startswith("answer:"):
|
||||||
|
answer = line[len("answer:") :].strip()
|
||||||
|
elif difficulty is None and lower.startswith("difficulty:"):
|
||||||
|
difficulty = line[len("difficulty:") :].strip()
|
||||||
|
|
||||||
|
if question and answer and difficulty:
|
||||||
|
return question, answer, difficulty
|
||||||
|
if len(lines) == 3:
|
||||||
|
return lines[0], lines[1], lines[2]
|
||||||
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def parse_multiple_qa_output(output: str) -> List[Tuple[str, str, str]]:
|
||||||
|
"""
|
||||||
|
Splits the output into blocks (separated by one or more blank lines) and
|
||||||
|
attempts to parse each as a QA pair.
|
||||||
|
|
||||||
|
Returns a list of successfully parsed QA tuples.
|
||||||
|
"""
|
||||||
|
blocks = re.split(r"\n\s*\n", output.strip())
|
||||||
|
qa_pairs = []
|
||||||
|
for block in blocks:
|
||||||
|
parsed = parse_qa_block(block)
|
||||||
|
if parsed:
|
||||||
|
qa_pairs.append(parsed)
|
||||||
|
return qa_pairs
|
||||||
|
|
||||||
|
|
||||||
|
def generate_question_batch_for_chunks(
|
||||||
|
chunks: List, num_questions: int = 2, difficulty: Optional[str] = None
|
||||||
|
) -> List[Dict]:
|
||||||
|
"""
|
||||||
|
Generates QA pairs for multiple chunks in batch.
|
||||||
|
|
||||||
|
For each chunk, generates questions based on its content only.
|
||||||
|
Each prompt instructs the model to output exactly three lines per QA pair with markers.
|
||||||
|
Failed prompts are retried once in batch; if still unsuccessful, they are skipped.
|
||||||
|
|
||||||
|
Returns a list of dicts with keys: "chunk_id", "question", "answer", "difficulty", "chunk_content".
|
||||||
|
"""
|
||||||
|
prompts = []
|
||||||
|
chunk_ids = []
|
||||||
|
chunk_contents = []
|
||||||
|
|
||||||
|
# Prepare prompts for each chunk
|
||||||
|
for i, chunk in enumerate(chunks):
|
||||||
|
current = chunk.page_content
|
||||||
|
prompt = (
|
||||||
|
f"You are a question generator. Generate {num_questions} questions based on the following text.\n"
|
||||||
|
"Rules:\n"
|
||||||
|
"1. Questions must be answerable using ONLY the information in the text\n"
|
||||||
|
"2. Answers must be directly stated in the text\n"
|
||||||
|
"3. Each question should test understanding of a different aspect of the text\n"
|
||||||
|
"4. Questions should be clear and specific\n"
|
||||||
|
"5. Answers should be concise and factual\n\n"
|
||||||
|
"For each QA pair, output exactly three lines with no extra commentary:\n"
|
||||||
|
"Line 1: Question: <your question>\n"
|
||||||
|
"Line 2: Answer: <the answer>\n"
|
||||||
|
"Line 3: Difficulty: <easy, medium, or hard>\n"
|
||||||
|
"Do not include any additional text.\n\n"
|
||||||
|
"Text:\n"
|
||||||
|
f"{current}\n"
|
||||||
|
)
|
||||||
|
prompts.append(prompt)
|
||||||
|
chunk_ids.append(i + 1) # 1-based indexing
|
||||||
|
chunk_contents.append(current)
|
||||||
|
|
||||||
|
# First batch generation
|
||||||
|
outputs = batch_generate(prompts)
|
||||||
|
results: List[Optional[List[Tuple[str, str, str]]]] = [None] * len(outputs)
|
||||||
|
failed_indices = []
|
||||||
|
|
||||||
|
# Parse each output
|
||||||
|
for idx, output in enumerate(outputs):
|
||||||
|
qa_pairs = parse_multiple_qa_output(output)
|
||||||
|
if qa_pairs is None or len(qa_pairs) < num_questions:
|
||||||
|
failed_indices.append(idx)
|
||||||
|
logger.warning(f"Failed to generate enough QA pairs for chunk {idx + 1}")
|
||||||
|
else:
|
||||||
|
# Validate that answers exist in chunk content
|
||||||
|
valid_pairs = []
|
||||||
|
for q, a, d in qa_pairs:
|
||||||
|
if a.lower() in chunk_contents[idx].lower():
|
||||||
|
valid_pairs.append((q, a, d))
|
||||||
|
else:
|
||||||
|
logger.warning(f"Answer not found in chunk content: {a}")
|
||||||
|
|
||||||
|
if len(valid_pairs) >= num_questions:
|
||||||
|
results[idx] = valid_pairs[:num_questions]
|
||||||
|
else:
|
||||||
|
failed_indices.append(idx)
|
||||||
|
logger.warning(f"Not enough valid QA pairs for chunk {idx + 1}")
|
||||||
|
|
||||||
|
# Retry failed prompts in batch
|
||||||
|
if failed_indices:
|
||||||
|
logger.info(f"Retrying {len(failed_indices)} failed prompt(s)...")
|
||||||
|
retry_prompts = [prompts[i] for i in failed_indices]
|
||||||
|
retry_outputs = batch_generate(retry_prompts)
|
||||||
|
for j, idx in enumerate(failed_indices):
|
||||||
|
qa_pairs = parse_multiple_qa_output(retry_outputs[j])
|
||||||
|
if qa_pairs is not None and len(qa_pairs) >= num_questions:
|
||||||
|
# Validate answers again
|
||||||
|
valid_pairs = []
|
||||||
|
for q, a, d in qa_pairs:
|
||||||
|
if a.lower() in chunk_contents[idx].lower():
|
||||||
|
valid_pairs.append((q, a, d))
|
||||||
|
|
||||||
|
if len(valid_pairs) >= num_questions:
|
||||||
|
results[idx] = valid_pairs[:num_questions]
|
||||||
|
else:
|
||||||
|
results[idx] = None
|
||||||
|
logger.warning(
|
||||||
|
f"Retry failed for chunk {idx + 1}: not enough valid QA pairs"
|
||||||
|
)
|
||||||
|
else:
|
||||||
|
results[idx] = None
|
||||||
|
logger.warning(f"Retry failed for chunk {idx + 1}: parsing failed")
|
||||||
|
|
||||||
|
# Build final output, skipping prompts that failed even after retry
|
||||||
|
final_questions = []
|
||||||
|
for i, qa_list in enumerate(results):
|
||||||
|
if qa_list is not None:
|
||||||
|
for qa in qa_list:
|
||||||
|
final_questions.append(
|
||||||
|
{
|
||||||
|
"chunk_id": chunk_ids[i],
|
||||||
|
"question": qa[0],
|
||||||
|
"answer": qa[1],
|
||||||
|
"difficulty": qa[2],
|
||||||
|
"chunk_content": chunk_contents[i],
|
||||||
|
}
|
||||||
|
)
|
||||||
|
|
||||||
|
logger.info(f"Generated {len(final_questions)} valid QA pairs")
|
||||||
|
return final_questions
|
||||||
|
|
||||||
|
|
||||||
|
# Generate QA pairs in batch (using a sliding window over the chunks)
|
||||||
|
all_questions = generate_question_batch_for_chunks(
|
||||||
|
chunks, num_questions=2, difficulty="medium"
|
||||||
|
)
|
||||||
|
print(f"Generated {len(all_questions)} QA pairs.")
|
||||||
|
|
||||||
|
# Save the QA pairs to a JSON file
|
||||||
|
questions_path = os.path.join("saved_data", "questions.json")
|
||||||
|
with open(questions_path, "w") as f:
|
||||||
|
json.dump(all_questions, f, indent=2)
|
||||||
|
print(f"Saved questions to {questions_path}")
|
||||||
@ -0,0 +1,113 @@
|
|||||||
|
{
|
||||||
|
"cells": [
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"!pip install matplotlib -q"
|
||||||
|
]
|
||||||
|
},
|
||||||
|
{
|
||||||
|
"cell_type": "code",
|
||||||
|
"execution_count": null,
|
||||||
|
"metadata": {},
|
||||||
|
"outputs": [],
|
||||||
|
"source": [
|
||||||
|
"import numpy as np\n",
|
||||||
|
"import matplotlib.pyplot as plt\n",
|
||||||
|
"\n",
|
||||||
|
"def plot_reward_functions():\n",
|
||||||
|
" # Generate retry counts from 0 to 15\n",
|
||||||
|
" retries = np.linspace(0, 15, 100)\n",
|
||||||
|
" \n",
|
||||||
|
" # 1. Basic Sigmoid\n",
|
||||||
|
" basic_sigmoid = 1 / (1 + np.exp(-(retries - 4)))\n",
|
||||||
|
" \n",
|
||||||
|
" # 2. Our Modified Sigmoid\n",
|
||||||
|
" x = retries - 4 # Center at 4 retries\n",
|
||||||
|
" modified_sigmoid = 1 / (1 + np.exp(-x + abs(x)/2))\n",
|
||||||
|
" \n",
|
||||||
|
" # 3. With Penalty\n",
|
||||||
|
" penalized_reward = modified_sigmoid.copy()\n",
|
||||||
|
" for i, r in enumerate(retries):\n",
|
||||||
|
" if r > 6:\n",
|
||||||
|
" penalty = 0.2 * (r - 6)\n",
|
||||||
|
" penalized_reward[i] = max(0.1, modified_sigmoid[i] - penalty)\n",
|
||||||
|
" \n",
|
||||||
|
" # Plotting\n",
|
||||||
|
" plt.figure(figsize=(12, 6))\n",
|
||||||
|
" \n",
|
||||||
|
" plt.plot(retries, basic_sigmoid, 'b--', label='Basic Sigmoid')\n",
|
||||||
|
" plt.plot(retries, modified_sigmoid, 'g--', label='Modified Sigmoid')\n",
|
||||||
|
" plt.plot(retries, penalized_reward, 'r-', label='Final Reward (with penalty)', linewidth=2)\n",
|
||||||
|
" \n",
|
||||||
|
" # Add vertical lines for key points\n",
|
||||||
|
" plt.axvline(x=4, color='gray', linestyle=':', alpha=0.5, label='Peak (4 retries)')\n",
|
||||||
|
" plt.axvline(x=6, color='gray', linestyle=':', alpha=0.5, label='Penalty Start (6 retries)')\n",
|
||||||
|
" \n",
|
||||||
|
" plt.grid(True, alpha=0.3)\n",
|
||||||
|
" plt.xlabel('Number of Retries')\n",
|
||||||
|
" plt.ylabel('Reward')\n",
|
||||||
|
" plt.title('Reward Function Visualization')\n",
|
||||||
|
" plt.legend()\n",
|
||||||
|
" plt.ylim(-0.1, 1.1)\n",
|
||||||
|
" \n",
|
||||||
|
" # Add annotations\n",
|
||||||
|
" plt.annotate('Optimal Zone', xy=(4, 0.8), xytext=(4, 0.9),\n",
|
||||||
|
" ha='center', va='bottom',\n",
|
||||||
|
" bbox=dict(boxstyle='round,pad=0.5', fc='yellow', alpha=0.3),\n",
|
||||||
|
" arrowprops=dict(arrowstyle='->'))\n",
|
||||||
|
" \n",
|
||||||
|
" plt.annotate('Penalty Zone', xy=(8, 0.3), xytext=(8, 0.5),\n",
|
||||||
|
" ha='center', va='bottom',\n",
|
||||||
|
" bbox=dict(boxstyle='round,pad=0.5', fc='red', alpha=0.3),\n",
|
||||||
|
" arrowprops=dict(arrowstyle='->'))\n",
|
||||||
|
" \n",
|
||||||
|
" plt.show()\n",
|
||||||
|
"\n",
|
||||||
|
"# Run the visualization\n",
|
||||||
|
"plot_reward_functions()\n",
|
||||||
|
"\n",
|
||||||
|
"# Print reward values for specific retry counts\n",
|
||||||
|
"def print_reward_examples():\n",
|
||||||
|
" retry_examples = [1, 2, 3, 4, 5, 6, 7, 8, 10, 12]\n",
|
||||||
|
" print(\"\\nReward values for different retry counts:\")\n",
|
||||||
|
" print(\"Retries | Reward\")\n",
|
||||||
|
" print(\"-\" * 20)\n",
|
||||||
|
" \n",
|
||||||
|
" for retries in retry_examples:\n",
|
||||||
|
" x = retries - 4\n",
|
||||||
|
" reward = 1 / (1 + np.exp(-x + abs(x)/2))\n",
|
||||||
|
" if retries > 6:\n",
|
||||||
|
" penalty = 0.2 * (retries - 6)\n",
|
||||||
|
" reward = max(0.1, reward - penalty)\n",
|
||||||
|
" print(f\"{retries:7d} | {reward:.3f}\")\n",
|
||||||
|
"\n",
|
||||||
|
"print_reward_examples()"
|
||||||
|
]
|
||||||
|
}
|
||||||
|
],
|
||||||
|
"metadata": {
|
||||||
|
"kernelspec": {
|
||||||
|
"display_name": ".venv",
|
||||||
|
"language": "python",
|
||||||
|
"name": "python3"
|
||||||
|
},
|
||||||
|
"language_info": {
|
||||||
|
"codemirror_mode": {
|
||||||
|
"name": "ipython",
|
||||||
|
"version": 3
|
||||||
|
},
|
||||||
|
"file_extension": ".py",
|
||||||
|
"mimetype": "text/x-python",
|
||||||
|
"name": "python",
|
||||||
|
"nbconvert_exporter": "python",
|
||||||
|
"pygments_lexer": "ipython3",
|
||||||
|
"version": "3.11.11"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
"nbformat": 4,
|
||||||
|
"nbformat_minor": 2
|
||||||
|
}
|
||||||
Loading…
Reference in new issue