Upload score_pilot_mdna_v7_hf_jobs.py with huggingface_hub

score_pilot_mdna_v7_hf_jobs.py

@@ -0,0 +1,431 @@
+# /// script
+# dependencies = [
+#   "transformers",
+#   "accelerate",
+#   "peft",
+#   "bitsandbytes",
+#   "sentencepiece",
+#   "huggingface_hub",
+#   "pandas",
+#   "tqdm",
+#   "scikit-learn"
+# ]
+# ///
+
+import os
+import json
+import gc
+from datetime import datetime
+
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+
+from huggingface_hub import hf_hub_download, HfApi
+from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
+from peft import PeftModel
+
+
+# =========================
+# 1. Basic configuration
+# =========================
+
+HF_TOKEN = os.environ.get("HF_TOKEN")
+
+if HF_TOKEN is None:
+    raise ValueError("HF_TOKEN is missing. Please pass it with --secrets HF_TOKEN.")
+
+MODEL_REPO = "AnneQtao/deepseek-r1-finsent-lora-v7"
+BASE_MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Llama-8B"
+
+INPUT_CSV = "pilot_mdna_2022_2024_30.csv"
+
+OUTPUT_CSV = "pilot_mdna_v7_full_document_scores_a100.csv"
+OUTPUT_METADATA = "pilot_mdna_v7_full_document_scores_a100_metadata.json"
+
+MAX_INPUT_TOKENS = 32768
+
+LABEL_OPTIONS = {
+    "Positive": "positive",
+    "Negative": "negative",
+    "Neutral": "neutral"
+}
+
+DISPLAY_LABELS = ["Positive", "Negative", "Neutral"]
+
+
+# =========================
+# 2. Helper functions
+# =========================
+
+def make_prompt_full_mdna(mdna_text: str) -> str:
+    return (
+        "Classify the sentiment of the following MD&A disclosure as "
+        "positive, negative, or neutral based on expected implications for "
+        "firm performance, demand, costs, margins, risks, uncertainty, and outlook.\n\n"
+        f"MD&A:\n{mdna_text}\n\n"
+        "Sentiment:"
+    )
+
+
+def safe_cuda_empty_cache():
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+@torch.no_grad()
+def label_logprob_no_truncation(model, tokenizer, prompt: str, label_text: str):
+    full_token_count = len(tokenizer(prompt, add_special_tokens=False).input_ids)
+
+    if full_token_count > MAX_INPUT_TOKENS:
+        return None, full_token_count, True
+
+    prompt_ids = tokenizer(
+        prompt,
+        return_tensors="pt",
+        truncation=False
+    ).input_ids.to(model.device)
+
+    label_ids = tokenizer(
+        label_text,
+        add_special_tokens=False,
+        return_tensors="pt"
+    ).input_ids.to(model.device)
+
+    input_ids = torch.cat([prompt_ids, label_ids], dim=1)
+
+    outputs = model(input_ids=input_ids, use_cache=False)
+    logits = outputs.logits
+
+    log_probs = F.log_softmax(logits[:, :-1, :], dim=-1)
+
+    prompt_len = prompt_ids.shape[1]
+    total_logprob = 0.0
+
+    for j in range(label_ids.shape[1]):
+        token_id = label_ids[0, j]
+        position = prompt_len + j - 1
+        total_logprob += log_probs[0, position, token_id].item()
+
+    del prompt_ids, label_ids, input_ids, outputs, logits, log_probs
+    safe_cuda_empty_cache()
+
+    return total_logprob, full_token_count, False
+
+
+def score_one_mdna(model, tokenizer, mdna_text: str):
+    prompt = make_prompt_full_mdna(mdna_text)
+    full_token_count = len(tokenizer(prompt, add_special_tokens=False).input_ids)
+
+    if full_token_count > MAX_INPUT_TOKENS:
+        return {
+            "score_status": "SKIPPED_TOO_LONG",
+            "score_token_count": int(full_token_count),
+            "max_input_tokens": int(MAX_INPUT_TOKENS),
+            "truncated": False,
+            "too_long_for_current_limit": True,
+            "P_pos": None,
+            "P_neg": None,
+            "P_neu": None,
+            "S_DS": None,
+            "predicted_label": "SKIPPED_TOO_LONG"
+        }
+
+    try:
+        scores = []
+
+        for label in DISPLAY_LABELS:
+            score, tc, too_long = label_logprob_no_truncation(
+                model=model,
+                tokenizer=tokenizer,
+                prompt=prompt,
+                label_text=LABEL_OPTIONS[label],
+            )
+
+            if too_long:
+                return {
+                    "score_status": "SKIPPED_TOO_LONG",
+                    "score_token_count": int(tc),
+                    "max_input_tokens": int(MAX_INPUT_TOKENS),
+                    "truncated": False,
+                    "too_long_for_current_limit": True,
+                    "P_pos": None,
+                    "P_neg": None,
+                    "P_neu": None,
+                    "S_DS": None,
+                    "predicted_label": "SKIPPED_TOO_LONG"
+                }
+
+            scores.append(score)
+
+        scores = torch.tensor(scores, dtype=torch.float32)
+        probs = F.softmax(scores, dim=0).cpu().numpy()
+
+        return {
+            "score_status": "SUCCESS",
+            "score_token_count": int(full_token_count),
+            "max_input_tokens": int(MAX_INPUT_TOKENS),
+            "truncated": False,
+            "too_long_for_current_limit": False,
+            "P_pos": float(probs[0]),
+            "P_neg": float(probs[1]),
+            "P_neu": float(probs[2]),
+            "S_DS": float(probs[0] - probs[1]),
+            "predicted_label": DISPLAY_LABELS[int(probs.argmax())]
+        }
+
+    except RuntimeError as e:
+        message = str(e)
+
+        if "out of memory" in message.lower() or "cuda" in message.lower():
+            safe_cuda_empty_cache()
+
+            return {
+                "score_status": "CUDA_OOM",
+                "score_token_count": int(full_token_count),
+                "max_input_tokens": int(MAX_INPUT_TOKENS),
+                "truncated": False,
+                "too_long_for_current_limit": False,
+                "P_pos": None,
+                "P_neg": None,
+                "P_neu": None,
+                "S_DS": None,
+                "predicted_label": "CUDA_OOM"
+            }
+
+        raise e
+
+
+# =========================
+# 3. Start job
+# =========================
+
+print("=" * 100)
+print("Starting V7 full-document MD&A scoring job")
+print("Time:", datetime.utcnow().isoformat())
+print("Model repo:", MODEL_REPO)
+print("Base model:", BASE_MODEL_NAME)
+print("Max input tokens:", MAX_INPUT_TOKENS)
+print("=" * 100)
+
+print("CUDA available:", torch.cuda.is_available())
+if torch.cuda.is_available():
+    print("GPU:", torch.cuda.get_device_name(0))
+    print("GPU memory allocated:", torch.cuda.memory_allocated() / 1024**3, "GB")
+
+
+# =========================
+# 4. Download input CSV
+# =========================
+
+print("Downloading input CSV from Hugging Face repo...")
+
+input_path = hf_hub_download(
+    repo_id=MODEL_REPO,
+    filename=INPUT_CSV,
+    repo_type="model",
+    token=HF_TOKEN
+)
+
+pilot_mdna_df = pd.read_csv(input_path)
+
+print("Input shape:", pilot_mdna_df.shape)
+print("Input columns:", pilot_mdna_df.columns.tolist())
+
+
+# =========================
+# 5. Keep valid MD&A
+# =========================
+
+pilot_ok_df = pilot_mdna_df[
+    (pilot_mdna_df["extract_status"] == "OK") &
+    (pilot_mdna_df["mdna_text"].notna()) &
+    (pilot_mdna_df["mdna_char_len"] > 2000)
+].copy()
+
+print("Usable MD&A filings:", len(pilot_ok_df))
+
+if len(pilot_ok_df) == 0:
+    raise ValueError("No usable MD&A filings found.")
+
+
+# =========================
+# 6. Load tokenizer + model + adapter
+# =========================
+
+print("Loading tokenizer...")
+
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_REPO,
+    trust_remote_code=True,
+    token=HF_TOKEN
+)
+
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+# Avoid tokenizer warning from shorter stored tokenizer_config
+tokenizer.model_max_length = MAX_INPUT_TOKENS
+
+print("Loading base model in 4-bit...")
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_quant_type="nf4"
+)
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    BASE_MODEL_NAME,
+    quantization_config=bnb_config,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    attn_implementation="sdpa"
+)
+
+base_model.config.use_cache = False
+
+print("Loading V7 LoRA adapter...")
+
+model = PeftModel.from_pretrained(
+    base_model,
+    MODEL_REPO,
+    token=HF_TOKEN
+)
+
+model.eval()
+model.config.use_cache = False
+
+print("✅ Base model + V7 LoRA adapter loaded successfully!")
+
+
+# =========================
+# 7. Token count
+# =========================
+
+print("Counting tokens...")
+
+pilot_ok_df["token_count"] = pilot_ok_df["mdna_text"].apply(
+    lambda x: len(tokenizer(make_prompt_full_mdna(x), add_special_tokens=False).input_ids)
+)
+
+print("Token count summary:")
+print(pilot_ok_df["token_count"].describe())
+
+for limit in [4096, 8192, 16384, 32768, 65536]:
+    n_over = (pilot_ok_df["token_count"] > limit).sum()
+    print(f"Over {limit} tokens: {n_over} / {len(pilot_ok_df)}")
+
+
+# =========================
+# 8. Score each MD&A
+# =========================
+
+print("Scoring full MD&A documents...")
+
+score_rows = []
+
+for idx, row in tqdm(pilot_ok_df.iterrows(), total=len(pilot_ok_df)):
+    ticker = row.get("ticker", "UNKNOWN")
+    filing_date = row.get("filing_date", "UNKNOWN")
+    token_count = row.get("token_count", None)
+
+    print("-" * 100)
+    print(f"Scoring: {ticker} | {filing_date} | tokens={token_count}")
+
+    result = score_one_mdna(
+        model=model,
+        tokenizer=tokenizer,
+        mdna_text=row["mdna_text"]
+    )
+
+    print("Result:", result)
+
+    merged = row.to_dict()
+    merged.update(result)
+    score_rows.append(merged)
+
+    safe_cuda_empty_cache()
+
+
+pilot_scored_df = pd.DataFrame(score_rows)
+
+# Put important columns first
+important_cols = [
+    "pilot_group", "ticker", "company_name", "form", "filing_date", "report_date",
+    "extract_status", "mdna_char_len", "token_count", "score_token_count",
+    "max_input_tokens", "score_status", "truncated", "too_long_for_current_limit",
+    "P_pos", "P_neg", "P_neu", "S_DS", "predicted_label", "filing_url"
+]
+
+existing_important_cols = [c for c in important_cols if c in pilot_scored_df.columns]
+other_cols = [c for c in pilot_scored_df.columns if c not in existing_important_cols]
+
+pilot_scored_df = pilot_scored_df[existing_important_cols + other_cols]
+
+pilot_scored_df.to_csv(OUTPUT_CSV, index=False)
+
+print("Saved output:", OUTPUT_CSV)
+print(pilot_scored_df[[
+    "pilot_group", "ticker", "form", "filing_date",
+    "token_count", "score_status",
+    "P_pos", "P_neg", "P_neu", "S_DS", "predicted_label"
+]])
+
+
+# =========================
+# 9. Save metadata
+# =========================
+
+metadata = {
+    "run_time_utc": datetime.utcnow().isoformat(),
+    "model_repo": MODEL_REPO,
+    "base_model": BASE_MODEL_NAME,
+    "input_csv": INPUT_CSV,
+    "output_csv": OUTPUT_CSV,
+    "max_input_tokens": MAX_INPUT_TOKENS,
+    "n_input_rows": int(len(pilot_mdna_df)),
+    "n_usable_mdna": int(len(pilot_ok_df)),
+    "n_success": int((pilot_scored_df["score_status"] == "SUCCESS").sum()),
+    "n_skipped_too_long": int((pilot_scored_df["score_status"] == "SKIPPED_TOO_LONG").sum()),
+    "n_cuda_oom": int((pilot_scored_df["score_status"] == "CUDA_OOM").sum()),
+    "token_count_summary": pilot_ok_df["token_count"].describe().to_dict(),
+}
+
+with open(OUTPUT_METADATA, "w") as f:
+    json.dump(metadata, f, indent=2)
+
+print("Saved metadata:", OUTPUT_METADATA)
+print(json.dumps(metadata, indent=2))
+
+
+# =========================
+# 10. Upload outputs to Hugging Face repo
+# =========================
+
+print("Uploading outputs to Hugging Face repo...")
+
+api = HfApi(token=HF_TOKEN)
+
+api.upload_file(
+    path_or_fileobj=OUTPUT_CSV,
+    path_in_repo=OUTPUT_CSV,
+    repo_id=MODEL_REPO,
+    repo_type="model"
+)
+
+api.upload_file(
+    path_or_fileobj=OUTPUT_METADATA,
+    path_in_repo=OUTPUT_METADATA,
+    repo_id=MODEL_REPO,
+    repo_type="model"
+)
+
+print("✅ Uploaded outputs to:", MODEL_REPO)
+print("Output CSV:", OUTPUT_CSV)
+print("Metadata:", OUTPUT_METADATA)
+print("Job finished.")