import argparse import os if os.environ.get("FORCE_CLI_TQDM") == "1": from tqdm import tqdm else: from tqdm.autonotebook import tqdm import torch from datasets import load_dataset from offline.text_processor import TextProcessor from offline.triplets_extractor import RebelExtractor from offline.triplets_validator import DeBERTa_Validator def ingest(override_args: argparse.Namespace): def process_text(text_chunk: str, sample_idx: int): """ Pipeline to produce triplets from a paragraph """ samplebar.set_description_str("[Splitting sentences...]") sentences = text_processor.split_sentences(text_chunk) if len(sentences) < 1: return [] rebel_inputs = [sentences[0]] for i in range(1, len(sentences)): rebel_inputs.extend([ sentences[i-1] + " " + sentences[i], sentences[i] ]) # Utilize REBEL seq2seq model to extract triplets triplets = [] samplebar.set_description_str("[Extracting triplets...]") for sent_idx, s in enumerate(tqdm(rebel_inputs, position=1, desc="Samples", leave=False, disable=os.environ.get("TQDM_DISABLE") == "1")): try: parsed = triplets_extractor.parse_output(triplets_extractor.extract(s)) triplets.extend(parsed) except Exception as e: tqdm.write(f"[WARN] Failed on sentence no. {sent_idx+1} of sample no. {sample_idx+1}: {e}") continue triplets = list(set(triplets)) if len(triplets) < 1: return [] # Utilize DeBERTa to validate triplets using the text as context samplebar.set_description_str("[Validating triplets with DeBERTa...]") valid_triplets = [] try: valid_triplets = triplets_validator.validate_triplets(triplets, text_chunk) except Exception as e: tqdm.write(f"[WARN] Validation of sample no. {sent_idx+1} failed: {e}") return [] # Return empty instead of un-validated facts # Apply normalization rules (lower-case, replace Greek characters, no weird spaces,...) # The use spaCy to lemmatize entites and simplify relation phrases refined_triplets = [] for s, r, o in valid_triplets: try: p_s = text_processor.normalize_text(text_processor.lemmatize(s)) p_r = text_processor.normalize_text(text_processor.lemmatize(text_processor.simplify_phrase(r))) p_o = text_processor.normalize_text(text_processor.lemmatize(o)) refined_triplets.append((p_s, p_r, p_o)) except: tqdm.write(f"[WARN] Triplets refining failed: ({s}, {r}, {o})") continue return refined_triplets # Allowed categories allowed_categories = ['core_clinical', 'basic_biology', 'pharmacology', 'psychiatry'] default_args = argparse.Namespace( data_path="cogbuji/medqa_corpus_en", category=allowed_categories, num_samples=100, start_idx=0, output_path='./artifacts/raw_triplets', cuda=False, cuda_extractor=False ) args = argparse.Namespace(**{**vars(default_args), **vars(override_args)}) # Safe-guard if args.start_idx < 0: print("Invalid starting index. Fallback to 0") args.start_idx = 0 args.category = eval(args.category) invalid_category = set(args.category) - set(allowed_categories) if len(invalid_category): print("Detecting invalid category:", ", ".join(invalid_category)) print("Allowed category:", ", ".join(allowed_categories)) return # Make directory for outputs os.makedirs(args.output_path, exist_ok=True) # GPU mode checker rebel_device = torch.device("cuda:0" if args.cuda_extractor and torch.cuda.is_available() else "cpu") reberta_device = torch.device("cuda:0" if args.cuda and torch.cuda.is_available() else "cpu") text_processor = TextProcessor() print("spaCy loaded.") triplets_extractor = RebelExtractor(device=rebel_device) print("REBEL loaded.") triplets_validator = DeBERTa_Validator(device=reberta_device) print("DeBERTa loaded.") for ctg in args.category: total_triplets = 0 print(f"Category: {ctg}") # Load dataset with the specified category dataset = load_dataset( args.data_path, name=ctg, split='train', trust_remote_code=True ) dataset_len = len(dataset) # Handle starting index out of index error if (args.start_idx >= dataset_len): tqdm.write(f"[WARN] Out of index for category [{ctg}]: starting index (start_idx) must be smaller than {dataset_len}. Skipping this category...") continue # Handle number of samples: automatically truncated to len(dataset) num_samples = dataset_len - args.start_idx if args.num_samples < num_samples: num_samples = args.num_samples else: tqdm.write(f"[WARN] Out of index: Automatically reducing num_samples from {args.num_samples} reduced to {num_samples}") # Create progress bar samplebar = tqdm(range(num_samples), desc="Samples", position=0, disable=os.environ.get("TQDM_DISABLE") == "1") # Reset existing(if) files open(os.path.join(args.output_path, f'{ctg}_numsamples{args.num_samples}_start{args.start_idx}.csv'), 'w').close() for i in samplebar: # Free GPU every 20 epochs if (i+1) % 20 == 0: torch.cuda.empty_cache() # Replacement when bars are disabled if os.environ.get("TQDM_DISABLE") == "1": if (i+1) % 10 == 0: tqdm.write(f"Sample {i+1}/{num_samples}") try: # Call the pipeline triplets = process_text(dataset[args.start_idx + i]['text'], i+1) total_triplets += len(triplets) # Update information samplebar.set_postfix({"Total num triplets" : total_triplets }) # Append triplets to file with open(os.path.join(args.output_path, f'{ctg}_numsamples{args.num_samples}_start{args.start_idx}.csv'), 'a', encoding="utf-8") as f: for s, r, o in triplets: f.write(f'{s},{r},{o}\n') except Exception as e: tqdm.write(f"Sample no.{i+1} failed: {e}") continue torch.cuda.empty_cache() if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('--data_path', type=str, default="cogbuji/medqa_corpus_en", help="Hugging Face Repository path of the dataset, set to the local path if needed") parser.add_argument('--category', type=str, default="['core_clinical', 'basic_biology', 'pharmacology', 'psychiatry']", help="Specify which categories to be used. Allowed category: core_clinical, basic_biology, pharmacology, psychiatry") parser.add_argument('--cuda', action='store_true', help="Use GPU for LLM refining. Require CUDA-compiled llama-cpp-python") parser.add_argument('--cuda_extractor', action='store_true', help="Use GPU for triplets extractor (REBEL)") parser.add_argument('--num_samples', type=int, default=100, help="Number of samples in a category to be used") parser.add_argument('--output_path', type=str, default='./artifacts/raw_triplets', help="Output directory for triplets CSV files") parser.add_argument('--start_idx', type=int, default=0, help="Starting sample index") args = parser.parse_args() ingest(args)