bertose-affinose-training-code / code /contrastive /extract_resolved_glycans.py
supanthadey1's picture
Add BERTose and AFFINose training code release
1d6f391 verified
Raw
History Blame Contribute Delete
1.01 kB
#!/usr/bin/env python3
import pickle, json, os
BASE = "/work/ratul1/supantha/glycan-SD-VS/bert_training_v3/v3.1_cluster_training"
INPUT = f"{BASE}/bert_v5_bpe_topo/data/sequences_bpe_expanded.pkl"
OUTPUT = f"{BASE}/bert_v5.1_contrastive/data/fully_resolved_84k.pkl"
STATS = f"{BASE}/bert_v5.1_contrastive/data/extraction_stats.json"
def has_ambiguity(g):
if '?' in g.get('wurcs', ''):
return True
return any('?' in str(t) for t in g.get('tokens', []))
print("Loading IPA data...")
with open(INPUT, 'rb') as f:
all_data = pickle.load(f)
print(f"Loaded: {len(all_data):,}")
resolved = [g for g in all_data if not has_ambiguity(g)]
print(f"Fully resolved: {len(resolved):,}")
print(f"Still ambiguous: {len(all_data) - len(resolved):,}")
os.makedirs(os.path.dirname(OUTPUT), exist_ok=True)
with open(OUTPUT, 'wb') as f:
pickle.dump(resolved, f)
print(f"Saved to: {OUTPUT}")
with open(STATS, 'w') as f:
json.dump({"total": len(all_data), "resolved": len(resolved)}, f, indent=2)