Update continuous_time_embedding.py

continuous_time_embedding.py

@@ -2,10 +2,8 @@
 Companion module for loading the continuous-time NewsBERT model.
 
 This model is not a plain AutoModelForMaskedLM. It wraps a full
-fine-tuned BERT with a continuous sinusoidal time embedding injected at the
-input layer. You need this file (or an equivalent copy of these two classes)
-to load and query it; `AutoModelForMaskedLM.from_pretrained(...)` alone will
-not work.
+fine-tuned BERT with a continuous sinusoidal (Fourier) time embedding
+injected at the input layer. You need this file to load and query it.
 
 Usage:
     from continuous_time_embedding import load_continuous_time_model
@@ -20,6 +18,8 @@ import torch.nn as nn
 from transformers import AutoTokenizer, AutoModelForMaskedLM
 from huggingface_hub import snapshot_download
 
+BASE_MODEL_ID = "TextMachineProject/NewsBERT_1800-1920" 
+
 MIN_YEAR = 1800.0
 MAX_YEAR = 1920.0
 MIN_PERIOD_YEARS = 5.0
@@ -30,8 +30,7 @@ class ContinuousTimeEmbedding(nn.Module):
     """Sinusoidal (Fourier) features over normalized year, projected to
     hidden_size. Nearby years produce nearby embeddings by construction.
     Frequency band: lowest = 1 cycle over the full 1800-1920 span, highest =
-    1 cycle per MIN_PERIOD_YEARS (5 years) -- see training script / model
-    card for the reasoning."""
+    1 cycle per MIN_PERIOD_YEARS (5 years)."""
 
     def __init__(self, hidden_size, n_freqs=N_TIME_FREQS, min_year=MIN_YEAR, max_year=MAX_YEAR,
                  min_period_years=MIN_PERIOD_YEARS):
@@ -103,13 +102,25 @@ class ContinuousTimeBertForMLM(nn.Module):
 
 def load_continuous_time_model(repo_id_or_path, device=None, **kwargs):
     """Convenience loader. Works with either a local checkpoint directory or
-    a Hugging Face Hub repo id (downloads it locally first)."""
+    a Hugging Face Hub repo id (downloads it locally first).
+    """
     if os.path.isdir(repo_id_or_path):
         local_dir = repo_id_or_path
     else:
         local_dir = snapshot_download(repo_id_or_path)
 
-    tokenizer = AutoTokenizer.from_pretrained(os.path.join(local_dir, "full_model"))
+    tokenizer = None
+    for candidate in (os.path.join(local_dir, "full_model"), local_dir):
+        try:
+            tokenizer = AutoTokenizer.from_pretrained(candidate)
+            break
+        except Exception:
+            continue
+    if tokenizer is None:
+        print(f"[load_continuous_time_model] No tokenizer found in {repo_id_or_path}, "
+              f"falling back to base model tokenizer: {BASE_MODEL_ID}")
+        tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_ID)
+
     model = ContinuousTimeBertForMLM.load_pretrained(local_dir, **kwargs)
 
     device = device or ("cuda" if torch.cuda.is_available() else "cpu")