File size: 2,980 Bytes
23a8f6d |
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 |
import argparse
import os
#!/usr/bin/env python3
import argparse
import os
import numpy as np
import pandas as pd
EPS = 1e-12
# ---------- Helper functions ----------
def logit(p: np.ndarray) -> np.ndarray:
"""Compute logit(p) = log(p / (1 - p))."""
p = np.clip(p, EPS, 1.0 - EPS)
return np.log(p / (1.0 - p))
def entropy(p: np.ndarray) -> np.ndarray:
"""Shannon entropy: H(p) = -Σ p log p."""
p = np.clip(p, EPS, 1.0)
return -np.sum(p * np.log(p), axis=1)
def top2_margin(p: np.ndarray) -> np.ndarray:
"""Margin = top1(p) - top2(p)."""
s = np.sort(p, axis=1)
return s[:, -1] - s[:, -2]
# ---------- Main ----------
def main():
ap = argparse.ArgumentParser(description="Compute probability-derived features (logit, max prob, margin, entropy).")
ap.add_argument("--input", required=True,
help="Path to FinSent_*_raw_probs.csv from FinBERT/RoBERTa step.")
ap.add_argument("--out_file", default=None,
help="Output CSV (default: adds _prob_features to filename).")
ap.add_argument("--out_dir", default="outputs", help="Base output directory")
ap.add_argument("--out_subdir", default="prob features", help="Subdirectory under out_dir to save prob features")
args = ap.parse_args()
df = pd.read_csv(args.input)
# Check probability columns
req = [
"fin_p_neg","fin_p_neu","fin_p_pos",
"rob_p_neg","rob_p_neu","rob_p_pos"
]
missing = [c for c in req if c not in df.columns]
if missing:
raise ValueError(f"Missing columns: {missing}")
# Prepare arrays
p_fin = df[["fin_p_neg","fin_p_neu","fin_p_pos"]].to_numpy(dtype=float)
p_rob = df[["rob_p_neg","rob_p_neu","rob_p_pos"]].to_numpy(dtype=float)
# ---- FinBERT features ----
fin_logit = logit(p_fin)
for i, cls in enumerate(["neg","neu","pos"]):
df[f"fin_logit_{cls}"] = fin_logit[:, i]
df["fin_max_prob"] = np.max(p_fin, axis=1)
df["fin_margin"] = top2_margin(p_fin)
df["fin_entropy"] = entropy(p_fin)
# ---- RoBERTa features ----
rob_logit = logit(p_rob)
for i, cls in enumerate(["neg","neu","pos"]):
df[f"rob_logit_{cls}"] = rob_logit[:, i]
df["rob_max_prob"] = np.max(p_rob, axis=1)
df["rob_margin"] = top2_margin(p_rob)
df["rob_entropy"] = entropy(p_rob)
# Save
root, ext = os.path.splitext(args.input)
if args.out_file:
out_path = args.out_file
else:
# save into outputs/<out_subdir>/ by default
save_dir = os.path.join(args.out_dir, args.out_subdir) if args.out_subdir else args.out_dir
os.makedirs(save_dir, exist_ok=True)
base = os.path.basename(root)
out_path = os.path.join(save_dir, f"{base}_prob_features.csv")
df.to_csv(out_path, index=False)
print(f"[✓] Saved probability-derived features to: {out_path}")
print("Added columns: fin/rob logits, max_prob, margin, entropy")
if __name__ == "__main__":
main() |