Duplicate from dejanseo/DEJAN-LM

Co-authored-by: Dan Petrovic <dejanseo@users.noreply.huggingface.co>

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,5 @@
+---
+license: other
+license_name: link-attribution
+license_link: https://dejanmarketing.com/link-attribution/
+---

app.py

@@ -0,0 +1,357 @@
+# app_interactive.py
+import streamlit as st
+import torch
+import random
+import os
+import pandas as pd
+from transformers import RobertaForMaskedLM, PreTrainedTokenizerFast
+import re
+
+# --- Configuration ---
+CHECKPOINT_BASE_DIR = "./checkpoints"
+PRESET_SENTENCE = "The quick brown fox jumps over the lazy dog near the river bank."
+TOP_K = 5
+
+# --- Initialize Session State ---
+if 'masked_indices' not in st.session_state:
+    st.session_state.masked_indices = set()
+if 'tokens' not in st.session_state:
+    st.session_state.tokens = []
+if 'token_ids' not in st.session_state:
+    st.session_state.token_ids = []
+if 'input_sentence' not in st.session_state:
+    st.session_state.input_sentence = PRESET_SENTENCE
+if 'display_tokens' not in st.session_state:
+    st.session_state.display_tokens = []
+
+# --- Helper Functions ---
+def sanitize_token_display(token):
+    """Clean up token display by removing special characters like Ġ."""
+    # Replace the 'Ġ' character with a more readable indicator
+    if isinstance(token, str) and token.startswith('Ġ'):
+        return token[1:]  # Remove the Ġ character
+    # Handle other special tokens if needed
+    elif token in ['<s>', '</s>', '<pad>']:
+        return token
+    else:
+        return token
+
+def find_checkpoints(base_dir):
+    """Finds valid checkpoint directories within the base directory."""
+    checkpoints = []
+    if not os.path.isdir(base_dir):
+        return checkpoints
+    for item in os.listdir(base_dir):
+        path = os.path.join(base_dir, item)
+        if os.path.isdir(path) and item.startswith("checkpoint-"):
+            if os.path.exists(os.path.join(path, "pytorch_model.bin")) or \
+               os.path.exists(os.path.join(path, "model.safetensors")):
+                 checkpoints.append(item)
+    checkpoints.sort(key=lambda x: int(re.search(r'(\d+)', x).group(1)))
+    return checkpoints
+
+@st.cache_resource
+def load_model_and_tokenizer(checkpoint_name):
+    """Loads the model and tokenizer from the specified checkpoint directory name."""
+    checkpoint_path = os.path.join(CHECKPOINT_BASE_DIR, checkpoint_name)
+    if not os.path.isdir(checkpoint_path):
+        st.error(f"Checkpoint directory not found: {checkpoint_path}")
+        return None, None
+    try:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = RobertaForMaskedLM.from_pretrained(checkpoint_path).to(device)
+        tokenizer = PreTrainedTokenizerFast.from_pretrained(checkpoint_path)
+        model.eval()
+        #st.success(f"Loaded {checkpoint_name} on {device}")
+        return model, tokenizer, device
+    except Exception as e:
+        st.error(f"Error loading {checkpoint_name}: {e}")
+        return None, None, None
+
+def tokenize_text(text, tokenizer):
+    """Tokenize the input text and return tokens and their IDs."""
+    encoding = tokenizer(text, return_tensors="pt", add_special_tokens=True)
+    input_ids = encoding.input_ids[0].tolist()
+    
+    # Get individual tokens
+    tokens = []
+    for id in input_ids:
+        token = tokenizer.convert_ids_to_tokens(id)
+        tokens.append(token)
+    
+    return tokens, input_ids
+
+def toggle_token(index):
+    """Toggle a token's masked status."""
+    if index in st.session_state.masked_indices:
+        st.session_state.masked_indices.remove(index)
+    else:
+        st.session_state.masked_indices.add(index)
+
+def update_input_sentence():
+    """Update the input sentence and reset masked indices."""
+    st.session_state.input_sentence = st.session_state.input_text
+    st.session_state.masked_indices = set()
+    
+def get_predictions(model, tokenizer, device):
+    """Get predictions for masked tokens."""
+    if not st.session_state.masked_indices:
+        return None, None, None, None
+    
+    # Create a copy of the token IDs
+    masked_input_ids = st.session_state.token_ids.copy()
+    
+    # Apply masks
+    for idx in st.session_state.masked_indices:
+        masked_input_ids[idx] = tokenizer.mask_token_id
+    
+    # Convert to tensor
+    masked_input_tensor = torch.tensor([masked_input_ids]).to(device)
+    
+    # Get predictions
+    with torch.no_grad():
+        outputs = model(input_ids=masked_input_tensor)
+        logits = outputs.logits
+    
+    results = []
+    top1_predictions = {}
+    prediction_tokens = {}
+    original_token_ranks = {}
+    
+    for masked_index in st.session_state.masked_indices:
+        mask_logits = logits[0, masked_index, :]
+        probabilities = torch.softmax(mask_logits, dim=-1)
+        top_k_probs, top_k_indices = torch.topk(probabilities, TOP_K)
+        
+        # Save top-1 prediction for reconstruction
+        top1_id = top_k_indices[0].item()
+        top1_predictions[masked_index] = top1_id
+        
+        # Sanitize the token here
+        raw_token = tokenizer.convert_ids_to_tokens(top1_id)
+        prediction_tokens[masked_index] = sanitize_token_display(raw_token)
+        
+        original_token = st.session_state.tokens[masked_index]
+        original_id = st.session_state.token_ids[masked_index]
+        
+        # Check if original token is in top K predictions
+        original_token_in_top_k = False
+        original_token_rank = -1  # -1 means not in top K
+        
+        for rank, token_id in enumerate(top_k_indices.tolist()):
+            predicted_token = tokenizer.convert_ids_to_tokens(token_id)
+            if predicted_token.lower() == original_token.lower() or token_id == original_id:
+                original_token_in_top_k = True
+                original_token_rank = rank
+                break
+        
+        original_token_ranks[masked_index] = original_token_rank
+        
+        for rank, (prob, token_id) in enumerate(zip(top_k_probs.tolist(), top_k_indices.tolist())):
+            predicted_token = tokenizer.convert_ids_to_tokens(token_id)
+            # Sanitize the predicted token for the results table
+            clean_predicted_token = sanitize_token_display(predicted_token)
+            
+            # Case insensitive match
+            is_match = predicted_token.lower() == original_token.lower()
+            results.append({
+                "Masked Index": masked_index,
+                "Rank": rank + 1,
+                "Predicted Token": clean_predicted_token,  # Use sanitized token
+                "Original Token": sanitize_token_display(original_token),  # Sanitize original token
+                "Exact Match": is_match,
+                "Probability": f"{prob:.4f}"
+            })
+    
+    # Reconstruct the sentence using top-1 predictions
+    reconstructed_ids = masked_input_ids.copy()
+    for idx in st.session_state.masked_indices:
+        reconstructed_ids[idx] = top1_predictions[idx]
+    
+    reconstructed_text = tokenizer.decode(reconstructed_ids, skip_special_tokens=True)
+    
+    return results, reconstructed_text, prediction_tokens, original_token_ranks
+
+# --- Streamlit App Layout ---
+
+st.set_page_config(layout="wide", page_title="Interactive MLM Inference")
+
+# Custom CSS to prevent text wrapping in buttons
+st.markdown("""
+<style>
+    .stButton button {
+        white-space: nowrap;
+        overflow: hidden;
+        text-overflow: ellipsis;
+        min-width: 80px;
+    }
+</style>
+""", unsafe_allow_html=True)
+
+st.title("🧪 Interactive MLM Inference")
+
+# --- Checkpoint Selection ---
+available_checkpoints = find_checkpoints(CHECKPOINT_BASE_DIR)
+
+if not available_checkpoints:
+    st.error(f"No checkpoints found in '{CHECKPOINT_BASE_DIR}'. Please train a model first.")
+    st.stop()
+
+selected_checkpoint = st.selectbox(
+    "Select Checkpoint:",
+    available_checkpoints,
+    index=len(available_checkpoints) - 1
+)
+
+# --- Load Model ---
+if selected_checkpoint:
+    model, tokenizer, device = load_model_and_tokenizer(selected_checkpoint)
+else:
+    model, tokenizer, device = None, None, None
+
+# --- Interactive Inference Section ---
+st.divider()
+st.subheader("Interactive Token Masking")
+
+# 1. Original text area
+st.text_area(
+    "Input Sentence:", 
+    value=st.session_state.input_sentence,
+    key="input_text",
+    on_change=update_input_sentence,
+    height=100
+)
+
+if model and tokenizer and device:
+    # Tokenize the input text
+    st.session_state.tokens, st.session_state.token_ids = tokenize_text(
+        st.session_state.input_sentence, 
+        tokenizer
+    )
+    
+    # Create sanitized display tokens
+    st.session_state.display_tokens = [sanitize_token_display(token) for token in st.session_state.tokens]
+    
+    # 2. Interactive token display
+    st.subheader("Click on tokens to mask/unmask them:")
+    
+    # Group tokens into rows (adjust number as needed)
+    tokens_per_row = 12
+    
+    # Calculate how many rows we need
+    num_rows = (len(st.session_state.tokens) + tokens_per_row - 1) // tokens_per_row
+    
+    for row in range(num_rows):
+        # Create columns for this row
+        start_idx = row * tokens_per_row
+        end_idx = min(start_idx + tokens_per_row, len(st.session_state.tokens))
+        row_tokens = st.session_state.tokens[start_idx:end_idx]
+        
+        # Create equal-width columns
+        cols = st.columns(len(row_tokens))
+        
+        for j, col in enumerate(cols):
+            idx = start_idx + j
+            token = st.session_state.tokens[idx]
+            
+            # Skip special tokens for masking
+            is_special = token in [
+                tokenizer.cls_token, 
+                tokenizer.sep_token, 
+                tokenizer.pad_token
+            ]
+            
+            is_masked = idx in st.session_state.masked_indices
+            
+            # Create a button for each token
+            button_key = f"token_{idx}"
+            button_label = sanitize_token_display(token) if not is_masked else "[MASK]"
+            
+            if col.button(
+                button_label,
+                key=button_key,
+                disabled=is_special,
+                help=f"Token ID: {st.session_state.token_ids[idx]}"
+            ):
+                toggle_token(idx)
+                st.rerun()
+    
+    # 3. Prediction area
+    if st.session_state.masked_indices:
+        results, reconstructed_text, prediction_tokens, original_token_ranks = get_predictions(model, tokenizer, device)
+        
+        st.subheader("Predictions:")
+        st.markdown("**Reconstructed sentence with predictions:**")
+        
+        # Create HTML for highlighting predictions
+        html = "<div style='padding: 10px; border-radius: 5px; border: 1px solid #ccc;'>"
+        
+        # Use the original tokenization to match masked positions
+        for i, token in enumerate(st.session_state.tokens):
+            # Skip special tokens 
+            if token in [tokenizer.cls_token, tokenizer.sep_token, tokenizer.pad_token]:
+                continue
+                
+            if i in st.session_state.masked_indices:
+                # This was a masked token
+                original_token = sanitize_token_display(st.session_state.tokens[i])
+                predicted_token = prediction_tokens[i]  # This is already sanitized in get_predictions
+                original_rank = original_token_ranks[i]
+                
+                # Color based on original token's rank in predictions
+                if original_rank == 0:  # Rank 0 means it was the top prediction
+                    # Green for top prediction (rank 1)
+                    html += f"<span style='background-color: #c3e6cb; padding: 2px 4px; border-radius: 3px; margin: 0 2px;'>{predicted_token}</span>"
+                elif original_rank != -1:  # In top 5 but not top
+                    # Blue for in top 5 but not top
+                    html += f"<span style='background-color: #b8daff; padding: 2px 4px; border-radius: 3px; margin: 0 2px;'>{predicted_token}</span>"
+                else:  # Not in top 5
+                    # Red for not in top 5
+                    html += f"<span style='background-color: #f8d7da; padding: 2px 4px; border-radius: 3px; margin: 0 2px;'>{predicted_token}</span>"
+            else:
+                # Not a masked token, display normally
+                sanitized_token = sanitize_token_display(token)
+                html += f"{sanitized_token} "
+        
+        html += "</div>"
+        
+        # Display the highlighted text
+        st.markdown(html, unsafe_allow_html=True)
+        
+        # Show detailed predictions
+        st.markdown("**Top predictions for each masked token:**")
+        
+        for masked_idx in st.session_state.masked_indices:
+            original_token = st.session_state.tokens[masked_idx]
+            original_rank = original_token_ranks[masked_idx]
+            
+            # Create a note about whether the original token was in top predictions
+            if original_rank == 0:
+                rank_note = "✅ Original token was the top prediction"
+            elif original_rank != -1:
+                rank_note = f"ℹ️ Original token was prediction #{original_rank+1}"
+            else:
+                rank_note = "❌ Original token not in top 5 predictions"
+            
+            # Sanitize the token display
+            clean_original_token = sanitize_token_display(original_token)
+            st.markdown(f"**Token {clean_original_token} at position {masked_idx}** - {rank_note}")
+            
+            # The dataframe is already sanitized in the get_predictions function
+            df = pd.DataFrame([r for r in results if r['Masked Index'] == masked_idx])
+            df = df[["Rank", "Predicted Token", "Probability"]]
+            
+            # Highlight the row with the original token if it's in top 5
+            if original_rank != -1:
+                # Use pandas styler to highlight the row
+                styled_df = df.style.apply(lambda x: ['background-color: #c3e6cb' if i == original_rank else '' for i in range(len(x))], axis=0)
+                st.dataframe(styled_df, use_container_width=True)
+            else:
+                st.dataframe(df, use_container_width=True)
+    else:
+        st.info("Click on tokens above to mask them and see predictions.")
+else:
+    st.warning("Please select a valid checkpoint to enable interactive masking.")
+
+st.divider()
+st.caption("Interactive app for RoBERTa Masked Language Modeling.")

config.json

@@ -0,0 +1,26 @@
+{
+  "architectures": [
+    "RobertaForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "bos_token_id": 2,
+  "classifier_dropout": null,
+  "eos_token_id": 3,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 256,
+  "initializer_range": 0.02,
+  "intermediate_size": 1024,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 514,
+  "model_type": "roberta",
+  "num_attention_heads": 8,
+  "num_hidden_layers": 4,
+  "pad_token_id": 0,
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.50.3",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 32000
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4f87a907da662f0f7cd1c78bc7e116dc34bf7bc822bd88d0d8be318cb9b6c530
+size 46336400

optimizer.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3540ab6a715c86446c0fcc17747212409c242e583821d6556df5df779c3b4fbc
+size 92717818

rng_state.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:36bbf17e45bd87663cd98ff4d6027892aa4320c31d67540c8ee33c1d805a30c7
+size 14244

scheduler.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8a971a8dd2d90d918014f25aed8de35f62e388573fdd5a7706b6d6fe96f8fb76
+size 1064

special_tokens_map.json

@@ -0,0 +1,37 @@
+{
+  "cls_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "<mask>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,54 @@
+{
+  "add_prefix_space": true,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<unk>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "</s>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<mask>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "clean_up_tokenization_spaces": false,
+  "cls_token": "<s>",
+  "extra_special_tokens": {},
+  "mask_token": "<mask>",
+  "model_max_length": 512,
+  "pad_token": "<pad>",
+  "sep_token": "</s>",
+  "tokenizer_class": "PreTrainedTokenizer",
+  "unk_token": "<unk>"
+}

train.py

@@ -0,0 +1,246 @@
+# train_fixed_clean_keys_v2.py
+
+import os
+import math
+import random
+import torch
+import pandas as pd
+import numpy as np
+import streamlit as st
+import plotly.graph_objects as go
+from transformers import (
+    RobertaConfig, RobertaForMaskedLM, Trainer, TrainingArguments,
+    PreTrainedTokenizerFast, DataCollatorForLanguageModeling, TrainerCallback
+)
+# Import Value from datasets alongside others
+from datasets import load_dataset, Features, Sequence, Value
+
+# --- Streamlit setup ---
+st.set_page_config(layout="wide")
+
+# --- Constants ---
+TOKENIZER_DIR = "tokenizer" # Ensure this matches the one used in preprocessing
+DATA_PATH = "training_data.jsonl" # Ensure this is the output from sentence_aware_processor.py
+OUTPUT_DIR = "./checkpoints"
+VOCAB_SIZE = 32000
+MAX_LEN = 512
+BATCH_SIZE = 64
+EPOCHS = 50
+GRAD_ACC = 8
+LEARNING_RATE = 1e-3
+MLM_PROB = 0.15
+SEED = 42
+
+# --- Seed ---
+def set_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+set_seed(SEED)
+
+# --- Tokenizer ---
+if not os.path.exists(os.path.join(TOKENIZER_DIR, "tokenizer.json")):
+    st.error(f"Tokenizer not found in {TOKENIZER_DIR}")
+    st.stop()
+try:
+    tokenizer = PreTrainedTokenizerFast.from_pretrained(TOKENIZER_DIR)
+    tokenizer.model_max_length = MAX_LEN
+except Exception as e:
+    st.error(f"Error loading tokenizer from {TOKENIZER_DIR}: {e}")
+    st.stop()
+
+
+# --- Dataset ---
+# !!! MODIFIED: Updated Features definition to match the JSONL structure !!!
+features = Features({
+    'id': Value(dtype='int64'),         # Added id field
+    'input_ids': Sequence(Value(dtype='int32')),
+    'source': Value(dtype='string')      # Added source field
+})
+# (Error handling remains)
+try:
+    # Load the dataset using the updated features
+    dataset = load_dataset("json", data_files=DATA_PATH, features=features, split="train")
+    #st.success(f"Loaded dataset from {DATA_PATH} with columns: {dataset.column_names}")
+except Exception as e:
+    st.error(f"Failed to load dataset from {DATA_PATH}: {e}")
+    st.info(f"Ensure '{DATA_PATH}' exists and matches the features: {features}")
+    st.stop()
+
+# --- Add Attention Mask ---
+# This function remains the same, as it only needs 'input_ids'
+if 'attention_mask' not in dataset.column_names:
+    def add_attention_mask(example):
+         # The length is derived from the 'input_ids' field
+         example["attention_mask"] = [1] * len(example["input_ids"])
+         return example
+    dataset = dataset.map(add_attention_mask, num_proc=max(1, os.cpu_count() // 2))
+    #st.info("Added 'attention_mask' column.")
+
+# --- Collator ---
+# DataCollatorForLanguageModeling will automatically ignore extra columns like 'id' and 'source'
+collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer,
+    mlm=True,
+    mlm_probability=MLM_PROB
+)
+
+# --- Model ---
+# Model definition remains the same
+config = RobertaConfig(
+    vocab_size=VOCAB_SIZE,
+    hidden_size=256,
+    num_hidden_layers=4,
+    num_attention_heads=8,
+    intermediate_size=1024,
+    max_position_embeddings=MAX_LEN + 2,
+    pad_token_id=tokenizer.pad_token_id,
+    bos_token_id=tokenizer.cls_token_id,
+    eos_token_id=tokenizer.sep_token_id,
+)
+model = RobertaForMaskedLM(config=config)
+
+# --- UI State ---
+# UI setup remains the same
+log = {"step": [], "loss": [], "grad_norm": [], "perplexity": []}
+progress = st.empty()
+col1, col2 = st.columns(2)
+with col1:
+    chart1_placeholder = st.empty()
+    chart2_placeholder = st.empty()
+with col2:
+    chart3_placeholder = st.empty()
+    chart4_placeholder = st.empty()
+
+# --- Plotting Functions (Unchanged) ---
+def get_safe_range(values, pad_percent=0.1):
+    values = pd.Series(values).dropna()
+    if values.empty: return (0, 1)
+    if len(values) == 1: return (values.iloc[0] * 0.9, values.iloc[0] * 1.1)
+    numeric_values = pd.to_numeric(values, errors='coerce').dropna()
+    if numeric_values.empty: return (0, 1)
+    low, high = np.percentile(numeric_values, [2, 95])
+    pad = abs(high - low) * pad_percent
+    return max(0, low - pad), high + pad
+
+def forecast_plot(df):
+    if len(df) < 10: return go.Figure(layout_title_text="Loss Forecast (Need more data)")
+    x = pd.to_numeric(df["step"], errors='coerce').dropna().values
+    y = pd.to_numeric(df["loss"], errors='coerce').dropna().values
+    if len(x) < 2 or len(y) < 2 or len(x) != len(y):
+         return go.Figure(layout_title_text="Loss Forecast (Data error)")
+
+    forecast_x = np.linspace(x[0], x[-1] * 1.5, 300)
+    fig = go.Figure()
+    fig.add_trace(go.Scatter(x=x, y=y, mode='lines', name="Actual Loss"))
+
+    for percent, color in [(1, 'orange'), (10, 'green'), (50, 'red')]:
+        n = max(5, int(len(x) * percent / 100))
+        if len(x) >= n and n >= 2:
+            sub_x, sub_y = x[-n:], y[-n:]
+            try:
+                valid_indices = ~np.isnan(sub_x) & ~np.isnan(sub_y)
+                if np.sum(valid_indices) >= 2:
+                    m, b = np.polyfit(sub_x[valid_indices], sub_y[valid_indices], 1)
+                    y_fit = m * forecast_x + b
+                    fig.add_trace(go.Scatter(x=forecast_x, y=y_fit, name=f"{percent}% Trend", line=dict(dash='dot', color=color)))
+            except (np.linalg.LinAlgError, ValueError) as e:
+                print(f"Warning: Could not fit trend for {percent}%: {e}")
+
+    fig.update_layout(title="Loss Forecast", xaxis_title="Step", yaxis_title="Loss", legend_title_text='Trend % (Recent)')
+    return fig
+
+# --- Streamlit Callback (Unchanged) ---
+class StreamlitCallback(TrainerCallback):
+    def on_log(self, args, state, control, logs=None, **kwargs):
+        if state.is_world_process_zero:
+            if logs is not None and "loss" in logs:
+                step = state.global_step
+                loss = float(logs["loss"]) if isinstance(logs["loss"], (int, float)) else None
+                grad = float(logs.get("grad_norm")) if isinstance(logs.get("grad_norm"), (int, float)) else None
+
+                if loss is not None:
+                    ppl = math.exp(min(loss, 700))
+                    log["step"].append(step)
+                    log["loss"].append(loss)
+                    log["grad_norm"].append(grad)
+                    log["perplexity"].append(ppl)
+
+                    df = pd.DataFrame(log).dropna(subset=['step', 'loss'])
+                    if not df.empty:
+                        try:
+                            r1 = get_safe_range(df["loss"])
+                            r2 = get_safe_range(df["grad_norm"])
+                            r3 = get_safe_range(df["perplexity"])
+
+                            fig1 = go.Figure().add_trace(go.Scatter(x=df["step"], y=df["loss"], mode='lines'))
+                            grad_norm_data = df["grad_norm"].dropna()
+                            if not grad_norm_data.empty:
+                                fig2 = go.Figure().add_trace(go.Scatter(x=df.loc[grad_norm_data.index, "step"], y=grad_norm_data, mode='lines'))
+                            else:
+                                fig2 = go.Figure()
+                            fig3 = go.Figure().add_trace(go.Scatter(x=df["step"], y=df["perplexity"], mode='lines'))
+
+                            fig1.update_layout(title="Loss", yaxis_range=r1, xaxis_title="Step", yaxis_title="Loss")
+                            fig2.update_layout(title="Gradient Norm", yaxis_range=r2, xaxis_title="Step", yaxis_title="Grad Norm")
+                            fig3.update_layout(title="Perplexity", yaxis_range=r3, xaxis_title="Step", yaxis_title="Perplexity")
+
+                            fig4 = forecast_plot(df)
+
+                            chart1_placeholder.plotly_chart(fig1, use_container_width=True, key=f"loss_chart_{step}")
+                            chart2_placeholder.plotly_chart(fig2, use_container_width=True, key=f"grad_norm_chart_{step}")
+                            chart3_placeholder.plotly_chart(fig3, use_container_width=True, key=f"perplexity_chart_{step}")
+                            chart4_placeholder.plotly_chart(fig4, use_container_width=True, key=f"forecast_chart_{step}")
+                        except Exception as e:
+                             print(f"Error updating Streamlit charts at step {step}: {e}")
+
+
+# --- Training args ---
+# Training args remain the same
+args = TrainingArguments(
+    output_dir=OUTPUT_DIR,
+    per_device_train_batch_size=BATCH_SIZE,
+    gradient_accumulation_steps=GRAD_ACC,
+    num_train_epochs=EPOCHS,
+    learning_rate=LEARNING_RATE,
+    lr_scheduler_type='linear',
+    warmup_ratio=0.1,
+    weight_decay=0.01,
+    max_grad_norm=1.0,
+    save_strategy="steps",
+    save_steps=1000,
+    save_total_limit=10,
+    logging_strategy="steps",
+    logging_steps=10,
+    dataloader_num_workers=4,
+    bf16=torch.cuda.is_bf16_supported(),
+    fp16=not torch.cuda.is_bf16_supported() and torch.cuda.is_available(),
+    seed=SEED,
+    report_to=["none"],
+    # !! Remember to handle checkpoints appropriately for a fresh run !!
+    resume_from_checkpoint=False, # Explicitly set to False for clean run
+)
+
+# --- Trainer ---
+# Trainer setup remains the same
+trainer = Trainer(
+    model=model,
+    args=args,
+    train_dataset=dataset,
+    data_collator=collator,
+    callbacks=[StreamlitCallback()]
+)
+
+# --- Train ---
+# Train call remains the same
+try:
+    # Start training (explicitly not resuming here due to args setting)
+    trainer.train() # No need to pass resume_from_checkpoint if set in args
+    progress.success("✅ Training complete.")
+    st.success("Training finished!")
+except Exception as e:
+    st.error(f"Training failed: {e}")
+    progress.error("❌ Training stopped due to error.")

train_tokenizer.py

@@ -0,0 +1,172 @@
+# improved_train_tokenizer_v2.py
+
+import os
+import sys
+from tokenizers import Tokenizer, models, pre_tokenizers, decoders, trainers, processors, normalizers
+from transformers import PreTrainedTokenizerFast
+
+# --- Configuration ---
+TRAIN_FILES = ["improved_sentences.txt"]  # Use the preprocessed file
+VOCAB_SIZE = 32000
+SPECIAL_TOKENS = ["<pad>", "<unk>", "<s>", "</s>", "<mask>"]
+OUTPUT_DIR = "./improved_tokenizer_v2"
+
+# --- Input File Check ---
+if not TRAIN_FILES or not os.path.exists(TRAIN_FILES[0]):
+    print(f"Error: Training file '{TRAIN_FILES[0]}' not found.")
+    sys.exit(1)
+
+print(f"Starting tokenizer training...")
+print(f"Training file(s): {TRAIN_FILES}")
+print(f"Target vocab size: {VOCAB_SIZE}")
+print(f"Output directory: {OUTPUT_DIR}")
+
+# --- Initialize Tokenizer ---
+# We'll use ByteLevel BPE with proper whitespace handling
+tokenizer = Tokenizer(models.BPE(unk_token="<unk>"))
+
+# --- Set Normalizer ---
+# This helps standardize the text before tokenization
+tokenizer.normalizer = normalizers.Sequence([
+    normalizers.NFC(),  # Unicode normalization
+    normalizers.Replace(r"\s+", " ")  # Replace multiple spaces with a single space
+])
+
+# --- Set Pre-tokenizer ---
+# This is critical for handling whitespace correctly
+tokenizer.pre_tokenizer = pre_tokenizers.ByteLevel(add_prefix_space=True)  # Back to True for proper space handling
+print(f"Using pre-tokenizer: ByteLevel(add_prefix_space=True)")
+
+# --- Set Decoder ---
+tokenizer.decoder = decoders.ByteLevel()
+print(f"Using decoder: {tokenizer.decoder.__class__.__name__}")
+
+# --- Define Trainer ---
+trainer = trainers.BpeTrainer(
+    vocab_size=VOCAB_SIZE,
+    special_tokens=SPECIAL_TOKENS,
+    show_progress=True,
+    initial_alphabet=pre_tokenizers.ByteLevel.alphabet(),
+)
+
+# --- Train Tokenizer ---
+print("\nTraining the tokenizer model (this might take a while)...")
+try:
+    tokenizer.train(files=TRAIN_FILES, trainer=trainer)
+    print("Training completed successfully.")
+except Exception as e:
+    print(f"\nError during tokenizer training: {e}")
+    sys.exit(1)
+
+# --- Add Post-processor ---
+tokenizer.post_processor = processors.TemplateProcessing(
+    single="<s> $A </s>",
+    pair="<s> $A </s> $B </s>",
+    special_tokens=[
+        ("<s>", tokenizer.token_to_id("<s>")),
+        ("</s>", tokenizer.token_to_id("</s>")),
+    ],
+)
+
+# --- Save Core Tokenizer ---
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+tokenizer_path = os.path.join(OUTPUT_DIR, "tokenizer.json")
+try:
+    tokenizer.save(tokenizer_path)
+    print(f"\nCore tokenizer saved to: {tokenizer_path}")
+except Exception as e:
+    print(f"Error saving core tokenizer: {e}")
+    sys.exit(1)
+
+# --- Create and Save HF Wrapper ---
+print("\nWrapping tokenizer with PreTrainedTokenizerFast...")
+try:
+    hf_tokenizer = PreTrainedTokenizerFast(
+        tokenizer_file=tokenizer_path,
+        unk_token="<unk>",
+        pad_token="<pad>",
+        cls_token="<s>",
+        sep_token="</s>",
+        mask_token="<mask>",
+        add_prefix_space=True  # Match the pre-tokenizer setting
+    )
+    hf_tokenizer.save_pretrained(OUTPUT_DIR)
+    print(f"Hugging Face compatible tokenizer files saved to: {OUTPUT_DIR}")
+except Exception as e:
+    print(f"Error saving Hugging Face tokenizer: {e}")
+    sys.exit(1)
+
+# --- Verification Step ---
+print("\n--- Verification ---")
+try:
+    print(f"Loading tokenizer for verification from: {OUTPUT_DIR}")
+    loaded_hf_tokenizer = PreTrainedTokenizerFast.from_pretrained(OUTPUT_DIR)
+
+    # Test multiple cases, especially those starting with periods or spaces
+    test_cases = [
+        "Simple sentence.",
+        " Sentence starting with space.",
+        "Sentence. Another sentence.",
+        ". Sentence starting with period.", 
+        "Word.Word",
+        "The quick brown fox jumps over the lazy dog."
+    ]
+    
+    print("\n=== Testing with new tokenizer ===")
+    for i, text in enumerate(test_cases):
+        print(f"\nTest {i+1}: '{text}'")
+        tokens = loaded_hf_tokenizer.tokenize(text)
+        print(f"Tokens: {tokens}")
+        
+        encoded = loaded_hf_tokenizer.encode(text, add_special_tokens=True)
+        decoded = loaded_hf_tokenizer.decode(encoded, skip_special_tokens=True)
+        print(f"Encoded: {encoded}")
+        print(f"Decoded: '{decoded}'")
+        
+        # Check if tokenization properly preserves content
+        if text.strip() == decoded.strip():
+            print("✓ Encoding/decoding preserved text content")
+        else:
+            print(f"⚠ Warning: Text content changed during encoding/decoding")
+            print(f"  Original: '{text}'")
+            print(f"  Decoded:  '{decoded}'")
+
+    # Check first token distributions
+    print("\n=== First Position Token Analysis ===")
+    print("Analyzing first token after <s> for potential bias...")
+    
+    # Simplified analysis of first token (just for demonstration)
+    from collections import Counter
+    first_token_counter = Counter()
+    
+    with open(TRAIN_FILES[0], 'r', encoding='utf-8') as f:
+        for i, line in enumerate(f):
+            if i >= 100:  # Just check first 100 lines
+                break
+            line = line.strip()
+            if not line:
+                continue
+                
+            encoded = loaded_hf_tokenizer.encode(line, add_special_tokens=True)
+            if len(encoded) > 1:  # Make sure there's at least one token after <s>
+                first_token_id = encoded[1]
+                first_token_counter[first_token_id] += 1
+    
+    total = sum(first_token_counter.values())
+    if total > 0:
+        print(f"\nTop 5 tokens at first position (after <s>) from {total} samples:")
+        for token_id, count in first_token_counter.most_common(5):
+            token_text = loaded_hf_tokenizer.decode([token_id])
+            percentage = (count / total) * 100
+            print(f"Token: '{token_text}' (ID: {token_id}) | Count: {count} | {percentage:.2f}%")
+            
+        # Specifically check period token
+        period_id = loaded_hf_tokenizer.encode('.', add_special_tokens=False)[0]
+        period_count = first_token_counter.get(period_id, 0)
+        period_percentage = (period_count / total) * 100 if total > 0 else 0
+        print(f"\nPeriod token ('.', ID: {period_id}) at first position: {period_count} times ({period_percentage:.2f}%)")
+    
+except Exception as e:
+    print(f"Error during verification: {e}")
+
+print("\n--- Tokenizer training script finished ---")

training_args.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b4255e8645ef79d5d89578e0550408329539962f036c0ac03649b791aa1cf604
+size 5304