File size: 5,708 Bytes
030876e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | import os
import json
import tqdm
import argparse
from openai import OpenAI
# -----------------------------
# CONFIGURATION
# -----------------------------
MODEL_NAME = "/home/mshahidul/readctrl_model/qwen3-32B_subclaims-extraction-8b_ctx"
API_URL = "http://localhost:8004/v1"
API_KEY = "EMPTY"
client = OpenAI(base_url=API_URL, api_key=API_KEY)
# -----------------------------
# SUBCLAIM EXTRACTION PROMPT
# -----------------------------
def extraction_prompt(medical_text: str) -> str:
return f"""
You are an expert medical annotator. Extract granular, factual subclaims.
A subclaim is the smallest standalone factual unit that can be independently verified.
Instructions:
1. Read the provided medical text.
2. Break it into clear, objective, atomic subclaims.
3. Each subclaim must come directly from the text.
4. Do not add, guess, or infer information.
5. Each subclaim should be short, specific, and verifiable.
6. Return ONLY a Python-style list of strings.
Medical Text:
{medical_text}
Return output as:
[
"subclaim 1",
"subclaim 2",
...
]
"""
# -----------------------------
# INFERENCE FUNCTION
# -----------------------------
def infer_subclaims(medical_text: str, temperature: float = 0.2) -> list:
if not medical_text or medical_text.strip() == "":
return []
final_prompt = extraction_prompt(medical_text)
try:
response = client.chat.completions.create(
model=MODEL_NAME,
messages=[{"role": "user", "content": final_prompt}],
max_tokens=1000,
temperature=temperature,
top_p=0.9,
)
res = response.choices[0].message.content.strip()
# Handle cases where the model might include <think> tags or markdown code blocks
if "</think>" in res:
res = res.split("</think>")[-1].strip()
if res.startswith("```json"):
res = res.replace("```json", "").replace("```", "").strip()
try:
return json.loads(res)
except:
# Fallback if JSON parsing fails but some text is returned
return [res]
except Exception as e:
print(f"API error for text snippet: {e}")
return []
# ... (Configuration and extraction_prompt remain the same) ...
# -----------------------------
# MAIN
# -----------------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--input_file", type=str,
default="/home/mshahidul/readctrl/data/synthetic_dataset_diff_labels/syn_data_with_gs_summary_en.json",
help="Path to input JSON file")
parser.add_argument("--start_index", type=int, default=0,
help="Start index for processing")
parser.add_argument("--end_index", type=int, default=-1,
help="End index for processing (exclusive). -1 = until end")
args = parser.parse_args()
SAVE_FOLDER = "/home/mshahidul/readctrl/data/extracting_subclaim"
os.makedirs(SAVE_FOLDER, exist_ok=True)
base_name = os.path.basename(args.input_file).replace(".json", "")
OUTPUT_FILE = os.path.join(
SAVE_FOLDER,
f"subclaims_with_generated_{base_name}_{args.start_index}_{args.end_index}.json"
)
print(f"Loading {args.input_file}...")
with open(args.input_file, "r") as f:
data = json.load(f)
total_items = len(data)
start = args.start_index
end = args.end_index if args.end_index != -1 else total_items
work_items = data[start:end]
result = []
if os.path.exists(OUTPUT_FILE):
try:
with open(OUTPUT_FILE, "r") as f:
result = json.load(f)
print(f"Resuming. {len(result)} items already processed.")
except:
result = []
# Using "index" or "id" as the unique identifier based on your JSON snippet
existing_ids = {r.get("index") or r.get("id") for r in result}
for item in tqdm.tqdm(work_items):
# Handle different ID key names
curr_id = item.get("index") if item.get("index") is not None else item.get("id")
if curr_id in existing_ids:
continue
# 1. Process standard fields
fulltext = item.get("fulltext", "")
summary = item.get("summary", "")
fulltext_sub = infer_subclaims(fulltext)
summary_sub = infer_subclaims(summary)
# 2. Process all generated texts (diff_label_texts)
# We will create a mirror dictionary to store the subclaims
diff_label_subclaims = {}
generated_texts = item.get("diff_label_texts", {})
for label, text in generated_texts.items():
if text:
diff_label_subclaims[label] = infer_subclaims(text)
else:
diff_label_subclaims[label] = []
# 3. Build output object
output_item = {
"index": curr_id,
"fulltext": fulltext,
"fulltext_subclaims": fulltext_sub,
"summary": summary,
"summary_subclaims": summary_sub,
"diff_label_texts": generated_texts,
"diff_label_subclaims": diff_label_subclaims, # New field
"readability_score": item.get("readability_score", None)
}
result.append(output_item)
# Periodic save
if len(result) % 10 == 0:
with open(OUTPUT_FILE, "w") as f:
json.dump(result, f, indent=4, ensure_ascii=False)
# Final save
with open(OUTPUT_FILE, "w") as f:
json.dump(result, f, indent=4, ensure_ascii=False)
print(f"Success! Results saved to: {OUTPUT_FILE}") |