# %%

import spacy

import coreferee

import torch

from transformers import (

AutoTokenizer,

AutoModelForTokenClassification,

AutoModelForSequenceClassification,

pipeline

)

import os

import json

import pickle

from datasets import load_dataset

from openai import OpenAI

import argparse

from tqdm import tqdm, trange

import networkx as nx

# %%

import logging

def setup_logging(log_file):

logger = logging.getLogger('my_logger')

logger.setLevel(logging.DEBUG)

file_handler = logging.FileHandler(log_file)

file_handler.setLevel(logging.DEBUG)

formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

file_handler.setFormatter(formatter)

logger.addHandler(file_handler)

return logger

def log_message(logger, message):

logger.info(message)

log_file = './data/app.log'

logger = setup_logging(log_file)

log_message(logger, 'For generating test cases')

argparser = argparse.ArgumentParser(description='Calculate the sum of two numbers')

argparser.add_argument('--ul_dataset', type=str, default='news', choices=['books', 'news'], help='unlearning dataset type')

# %%

## Load data

def load_json(file_path):

with open(file_path, 'r') as f:

data = json.load(f)

return data

raw_data = load_json(f'./data/muse_data/{argparser.ul_dataset}/raw/forget.json')

# %%

device='cuda'

nlp = spacy.load('en_core_web_lg')

nlp.add_pipe('coreferee')

triplet_extractor = pipeline('text2text-generation', model='Babelscape/rebel-large', max_length=512, tokenizer='Babelscape/rebel-large' ,device=device)

tokenizer = AutoTokenizer.from_pretrained("Babelscape/rebel-large")

# %%

def resolve_coreferences(text):

"""

Perform coreference resolution on input text using spaCy and coreferee.

Returns the text with resolved references.

"""

# Load spaCy model and add coreferee to pipeline

doc = nlp(text)

# Get all resolved chains

chains = doc._.coref_chains

# Store resolved text parts with whitespace

resolved_text = list(tok.text_with_ws for tok in doc)

# Print original coreference chains for analysis

# print("Coreference chains found:")

if chains:

for chain_index, chain in enumerate(chains):

# print(f"\nChain {chain_index}:")

for mention in chain:

mention_span = doc[mention[0]:mention[-1] + 1]

# print(f" - {mention_span.text}")

else:

# print("No coreference chains found.")

pass

# Resolve references with proper spacing

if chains:

for chain in chains:

# Get the first mention as the referent

first_mention = chain[0]

first_mention_span = doc[first_mention[0]:first_mention[-1] + 1]

# Replace later mentions with the first mention

for mention in chain[1:]:

if doc[mention[0]].text.lower() in ['he', 'she', 'it', 'they', 'his', 'her', 'its', 'their']:

# Replace the first token with referent and maintain its trailing whitespace

resolved_text[mention[0]] = first_mention_span.text + doc[mention[0]].whitespace_

# Clear other tokens in the mention if any, but maintain their whitespace

for token_idx in mention[1:]:

resolved_text[token_idx] = doc[token_idx].whitespace_

# Join all tokens to create final resolved text

resolved = ''.join(resolved_text)

return resolved

def process_text_with_spacy(text):

# 使用 spaCy 分句

doc = nlp(text)

sentences = list(doc.sents) # 分句处理

results = []

for sentence in sentences:

entities = [(ent.text, ent.label_, ent.start, ent.end) for ent in sentence.ents]

relations = extract_entity_relations(sentence)

results.append({

"sentence": sentence.text,

"entities": entities,

"relations": relations

})

return results

def extract_entity_relations(sentence):

"""Extract entity relations from a spaCy sentence."""

relations = []

entities = [ent for ent in sentence.ents]

# 遍历所有实体对

for i, ent1 in enumerate(entities):

for j, ent2 in enumerate(entities):

if i >= j:

continue

span_between = sentence[ent1.end:ent2.start] if ent1.end < ent2.start else None

if span_between and any(token.dep_ in ("prep", "pobj", "nsubj", "dobj") for token in span_between):

relations.append({

"entity1": (ent1.text, ent1.label_),

"entity2": (ent2.text, ent2.label_),

"context": span_between.text

})

return relations

def extract_triplets(text):

triplets = []

relation, subject, relation, object_ = '', '', '', ''

text = text.strip()

current = 'x'

for token in text.replace("<s>", "").replace("<pad>", "").replace("</s>", "").split():

if token == "<triplet>":

current = 't'

if relation != '':

triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})

relation = ''

subject = ''

elif token == "<subj>":

current = 's'

if relation != '':

triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})

object_ = ''

elif token == "<obj>":

current = 'o'

relation = ''

else:

if current == 't':

subject += ' ' + token

elif current == 's':

object_ += ' ' + token

elif current == 'o':

relation += ' ' + token

if subject != '' and relation != '' and object_ != '':

triplets.append({'head': subject.strip(), 'type': relation.strip(),'tail': object_.strip()})

return triplets

def single_IE_extractor(input_text):

extracted_text = triplet_extractor.tokenizer.batch_decode([triplet_extractor(input_text, return_tensors=True, return_text=False)[0]["generated_token_ids"]])

# Function to parse the generated text and extract the triplets

extracted_triplets = extract_triplets(extracted_text[0])

return extracted_triplets

def batch_IE_extractor(input_texts):

if isinstance(input_texts, str):

input_texts = [input_texts]

# 创建原始文本的索引列表

text_indices = list(range(len(input_texts)))

# 批量处理文本

extracted_texts = triplet_extractor.tokenizer.batch_decode(

[triplet_extractor(text, return_tensors=True, return_text=False)[0]["generated_token_ids"]

for text in input_texts]

)

# 提取三元组

extracted_triplets = [extract_triplets(text) for text in extracted_texts]

# 返回三元组结果和对应的原始文本索引

return extracted_triplets, text_indices

print(batch_IE_extractor(["The capital of France is Paris.","Alice like xiaoming."]))

def split_and_recombine(text, tokenizer, max_tokens):

"""

"""

# 用 spaCy 分句

doc = nlp(text)

sentences = [sent.text for sent in doc.sents]

return sentences

def parse_qa_response(llm_response: str) -> dict:

"""

Parse the LLM response and extract QA pairs.

Handles potential JSON formatting issues.

"""

try:

# Try direct JSON parsing first

if llm_response.startswith('''```'''):

llm_response.strip().strip("```").strip("json")

# print(llm_response)

qa_pairs = json.loads(llm_response)

# Validate the format

for key, value in qa_pairs.items():

if not isinstance(value, dict):

raise ValueError(f"Invalid format for question {key}")

if "question" not in value or "reference_answer" not in value:

raise ValueError(f"Missing required fields in question {key}")

return qa_pairs

except json.JSONDecodeError:

# Handle potential formatting issues

# You might want to add more sophisticated parsing here

log_message(logger, 'Invalid JSON format in LLM response')

raise ValueError("Invalid JSON format in LLM response")

client = OpenAI(api_key="sk-e9d569592d094342b271512258832ca8", base_url="https://api.deepseek.com")

def Query(sys, usr):

response = client.chat.completions.create(

model="deepseek-chat",

# model = "qwen-turbo",

# model = "gpt-4o",

messages=[

{"role": "system", "content": sys},

{"role": "user", "content": usr},

],

stream=False

)

return response.choices[0].message.content

# %%

# Prompt template

SYS_PROMPT = """You are an expert quiz generator. Given a text passage and a relationship triple, generate specific questions to test knowledge about this relationship based on the context provided.

Input Format:

- Text: A passage containing information about the relationship

- Relationship: A triple containing {'head': entity1, 'type': relation_type, 'tail': entity2}

Task:

Generate up to 5 focused questions that test understanding of the relationship between the head entity and tail entity, considering:

1. Questions should be answerable solely from the given context

2. Questions should be specific enough to have a unique correct answer

3. Questions can ask about the tail entity given the head entity and relationship type

4. Questions can ask about the relationship between the two entities

5. Questions can ask about specific details that establish this relationship

Requirements:

1. Each question must have a clear, unambiguous answer based on the context

2. Avoid overly broad or general questions

3. Focus on the specific relationship provided

4. Use the context to add specific details to questions

5. Ensure questions and answers are factually consistent with the provided text

Response Format:

The response must be a valid JSON object with the following structure:

{

"1": {

"question": "Your question text here",

"reference_answer": "The correct answer based on context"

},

"2": {

"question": "...",

"reference_answer": "..."

}

// ... up to 5 questions

}

Directly output the json, no anything else! NO ```json```！

Example Input:

Text: "The Greek Orthodox Church observes Lent as a period of fasting and spiritual reflection that begins on Clean Monday and lasts for 40 days. During this time, adherents follow strict dietary restrictions and increase their prayer and attendance at special services."

Relationship: {'head': 'Lent', 'type': 'religion', 'tail': 'Greek Orthodox'}

Example Output:

{

"1": {

"question": "Which religious denomination observes Lent beginning on Clean Monday with a 40-day period of fasting and spiritual reflection?",

"reference_answer": "Greek Orthodox"

},

"2": {

"question": "In the Greek Orthodox tradition, what is the length of the Lent period?",

"reference_answer": "40 days"

}

}

"""

USER_PROMPT = """

Please generate questions based on the following input:

Text: {text}

Relationship: {relationship}

"""

# %%

def add_edge_with_id(G, head, tail, relation):

# 初始化边的编号（如果图为空）

if not hasattr(G, 'next_edge_id'):

G.next_edge_id = 0

# 检查是否存在相同关系的边

found_edge_id = None

if G.has_edge(head, tail):

edges = G.get_edge_data(head, tail)

for key, data in edges.items():

if data['relation'] == relation:

# 找到相同关系的边,更新weight

G.remove_edge(head, tail, key)

new_weight = data.get('weight', 1) + 1

edge_id = data.get('edge_id')

G.add_edge(head, tail, edge_id=edge_id, relation=relation, weight=new_weight)

found_edge_id = edge_id

break

if found_edge_id is None:

# 没找到相同关系的边,创建新边

edge_id = G.next_edge_id

G.add_edge(head, tail, edge_id=edge_id, relation=relation, weight=1)

G.next_edge_id += 1

found_edge_id = edge_id

return found_edge_id

def IE_for_dataset(corpus, batch_size = 32):

G = nx.MultiDiGraph()

Edge2text = {}

Edge2corpu = {}

if not hasattr(G, 'next_edge_id'):

G.next_edge_id = 0

for idx in trange(len(corpus)):

# for idx in trange(10):

item = corpus[idx]

chunks = split_and_recombine(item, tokenizer, 512)

# slice the chunks into batches

for i in range(0, len(chunks), batch_size):

batch_chunks = chunks[i:min(len(chunks),i + batch_size)]

batch_results, _ = batch_IE_extractor(batch_chunks)

for i in range(len(batch_results)):

results = batch_results[i]

for ie in results:

head = ie['head']

tail = ie['tail']

relation = ie['type']

edge_id = add_edge_with_id(G, head, tail, relation)

if edge_id not in Edge2text:

Edge2text[edge_id] = []

if edge_id not in Edge2corpu:

Edge2corpu[edge_id] = []

Edge2corpu[edge_id].append(batch_chunks[i])

Edge2text[edge_id].append(batch_results[i])

return G, Edge2text, Edge2corpu

# %%

forget_data = load_json(f'./data/muse_data/{argparser.ul_dataset}/raw/forget.json')

retain_data = load_json(f'./data/muse_data/{argparser.ul_dataset}/raw/retain2.json')

# %%

G_forget, Edge2text_forget, Edge2corpu_forget = IE_for_dataset(forget_data)

G_retain, Edge2text_retain, Edge2corpu_retain = IE_for_dataset(retain_data)

with open(f'./data/{argparser.ul_dataset}_KG.pkl', 'wb') as f:

pickle.dump((G_forget, Edge2text_forget, Edge2corpu_forget, G_retain, Edge2text_retain, Edge2corpu_retain), f)

def topology_diff(G1, G2):

def get_edge_tuples(G):

edge_tuples = set()

for u, v in G.edges():

# 先获取边的数据

edge_data = G.get_edge_data(u, v)

# 然后再确保节点顺序一致

u_sorted, v_sorted = sorted([u, v])

for _, data in edge_data.items():

edge_tuples.add((u_sorted, v_sorted, data.get('relation')))

return edge_tuples

# 获取两个图的边集合（考虑relation）

edges1 = get_edge_tuples(G1)

edges2 = get_edge_tuples(G2)

# 计算差集

diff_edges = edges1 - edges2

# 创建新图，添加差集中的边（保留原图中的所有属性）

result = nx.MultiGraph()

for u, v, rel in diff_edges:

# 需要检查两个方向

if G1.has_edge(u, v):

edge_data = G1.get_edge_data(u, v)

else:

edge_data = G1.get_edge_data(v, u)

for edge_id, data in edge_data.items():

if data.get('relation') == rel:

data = data.copy() # 创建副本避免修改原图

data['original_edge_id'] = data.get('edge_id')

result.add_edge(u, v, **data)

return result

cleaned_graph = topology_diff(G_forget, G_retain)

with open(f'./data/{argparser.ul_dataset}_cleaned_graph.pkl', 'wb') as f:

pickle.dump(cleaned_graph, f)