File size: 6,893 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 176 177 178 179 180 | import os
# Set GPU environment variables
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
import torch
from unsloth import FastLanguageModel
import json
import tqdm
import argparse
# -----------------------------
# MODEL CACHE
# -----------------------------
_model_cache = {"model": None, "tokenizer": None}
def load_finetuned_model(model_path: str):
if _model_cache["model"] is not None:
return _model_cache["model"], _model_cache["tokenizer"]
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_path,
max_seq_length=8192,
load_in_4bit=False,
load_in_8bit=False,
full_finetuning=False,
)
_model_cache["model"], _model_cache["tokenizer"] = model, tokenizer
return model, tokenizer
# -----------------------------
# SUBCLAIM EXTRACTION PROMPT
# -----------------------------
def extraction_prompt(medical_text: str) -> str:
prompt = f"""
You are an expert medical annotator. Your task is to extract granular, factual subclaims from medical text.
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. Return ONLY a valid JSON list of strings.
Medical Text:
{medical_text}
Return your output in JSON list format:
[
"subclaim 1",
"subclaim 2"
]
"""
return prompt
# -----------------------------
# INFERENCE FUNCTION WITH AUTO-RETRY
# -----------------------------
def infer_subclaims(medical_text: str, model, tokenizer, temperature: float = 0.2, max_tokens: int = 2048, retries: int = 1) -> list:
if not medical_text or medical_text.strip() == "":
return []
prompt = extraction_prompt(medical_text)
messages = [{"role": "user", "content": prompt}]
chat_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(chat_text, return_tensors="pt").to("cuda")
with torch.no_grad():
output_ids = model.generate(
**inputs,
max_new_tokens=max_tokens,
temperature=temperature,
do_sample=False
)
output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
# Remove reasoning if model is a "Thinker" model
if "</think>" in output_text:
output_text = output_text.split("</think>")[-1].strip()
# JSON Parsing Logic
try:
start_idx = output_text.find('[')
end_idx = output_text.rfind(']') + 1
# Check if we have a complete bracketed pair
if start_idx != -1 and end_idx > start_idx:
content = output_text[start_idx:end_idx]
parsed = json.loads(content)
if isinstance(parsed, list):
return parsed
# If we are here, it means parsing failed or brackets were incomplete (truncation)
raise ValueError("Incomplete JSON list")
except (json.JSONDecodeError, ValueError):
# If truncation happened and we have retries left, double the tokens
if retries > 0:
new_max = max_tokens + 2048 # Increment by 2k tokens
print(f"\n[Warning] Truncation detected. Retrying with {new_max} tokens...")
return infer_subclaims(medical_text, model, tokenizer, temperature, max_tokens=new_max, retries=retries-1)
# Final fallback: return the raw text wrapped in a list so the pipeline doesn't crash
return [output_text]
# -----------------------------
# MAIN EXECUTION
# -----------------------------
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--input_file", type=str, required=True)
args = parser.parse_args()
INPUT_FILE = args.input_file
file_name = os.path.basename(INPUT_FILE).split(".json")[0]
SAVE_FOLDER = "/home/mshahidul/readctrl/data/extracting_subclaim"
MODEL_PATH = "/home/mshahidul/readctrl_model/qwen3-32B_subclaims-extraction-8b_ctx"
os.makedirs(SAVE_FOLDER, exist_ok=True)
OUTPUT_FILE = os.path.join(SAVE_FOLDER, f"extracted_subclaims_{file_name}.json")
model, tokenizer = load_finetuned_model(MODEL_PATH)
with open(INPUT_FILE, "r") as f:
data = json.load(f)
result = []
if os.path.exists(OUTPUT_FILE):
with open(OUTPUT_FILE, "r") as f:
result = json.load(f)
processed_data = {str(item.get("index") or item.get("id")): item for item in result}
for item in tqdm.tqdm(data):
item_id = str(item.get("index") if item.get("index") is not None else item.get("id"))
existing_entry = processed_data.get(item_id)
# 1. Process Fulltext (The longest field, high initial token count)
if not existing_entry or not isinstance(existing_entry.get("fulltext_subclaims"), list):
f_sub = infer_subclaims(item.get("fulltext", ""), model, tokenizer, max_tokens=3072, retries=2)
else:
f_sub = existing_entry["fulltext_subclaims"]
# 2. Process Summary
if not existing_entry or not isinstance(existing_entry.get("summary_subclaims"), list):
s_sub = infer_subclaims(item.get("summary", ""), model, tokenizer, max_tokens=2048, retries=1)
else:
s_sub = existing_entry["summary_subclaims"]
# 3. Process All Generated Texts (diff_label_texts)
diff_label_texts = item.get("diff_label_texts", {})
diff_label_subclaims = existing_entry.get("diff_label_subclaims", {}) if existing_entry else {}
for label, text in diff_label_texts.items():
if label not in diff_label_subclaims or not isinstance(diff_label_subclaims[label], list):
# Generated texts are shorter, but we still allow 1 retry
diff_label_subclaims[label] = infer_subclaims(text, model, tokenizer, max_tokens=1536, retries=1)
# 4. Save
new_entry = {
"index": item.get("index"),
"id": item.get("id"),
"fulltext": item.get("fulltext", ""),
"fulltext_subclaims": f_sub,
"summary": item.get("summary", ""),
"summary_subclaims": s_sub,
"diff_label_texts": diff_label_texts,
"diff_label_subclaims": diff_label_subclaims,
"readability_score": item.get("readability_score", None)
}
processed_data[item_id] = new_entry
if len(processed_data) % 10 == 0:
with open(OUTPUT_FILE, "w") as f:
json.dump(list(processed_data.values()), f, indent=4, ensure_ascii=False)
with open(OUTPUT_FILE, "w") as f:
json.dump(list(processed_data.values()), f, indent=4, ensure_ascii=False)
print(f"Extraction completed. File saved at: {OUTPUT_FILE}") |