Upload models.py with huggingface_hub

models.py

@@ -0,0 +1,258 @@
+from random import randint
+from string import printable
+import numpy as np
+import torch
+from rapidfuzz.distance.Levenshtein import distance as ldistance
+from torch.optim import AdamW
+from models import EditDistanceModel
+
+def pad_with_null(string: str, target_length: int):
+    null_char = "\0"
+    padding_needed = max(0, target_length - len(string))
+    return (string + (null_char * padding_needed))[:target_length]
+
+def string_to_tensor(string: str, length: int) -> torch.Tensor:
+    """Converts a string to a tensor of character indices."""
+    padded = pad_with_null(string, length)
+    # Use ord() to get integer representation, clamp to vocab size
+    indices = [min(ord(c), 127) for c in padded]
+    return torch.tensor(indices, dtype=torch.long)
+
+def random_char() -> str:
+    pos = randint(0, len(printable) - 1)
+    return printable[pos]
+
+def random_str(length: int) -> str:
+    return "".join([random_char() for _ in range(length)])
+
+def mangle_string(source: str, d: int) -> str:
+    """
+    Efficiently mangles a string to approximately the target distance
+    Uses list operations for better performance
+    """
+    if d <= 0:
+        return source
+
+    mangled = list(source)
+    edits_made = 0
+    max_attempts = d * 3  # Prevent infinite loops
+    attempts = 0
+
+    while edits_made < d and attempts < max_attempts:
+        attempts += 1
+
+        if len(mangled) == 0:
+            position = 0
+            edit = "insert"
+        else:
+            position = randint(0, len(mangled) - 1)
+            edit = ["insert", "delete", "modify"][randint(0, 2)]
+
+        if edit == "insert":
+            mangled.insert(position, random_char())
+            edits_made += 1
+        elif edit == "modify" and len(mangled) > 0:
+            old_char = mangled[position]
+            new_char = random_char()
+            if old_char != new_char:  # Only count as edit if actually different
+                mangled[position] = new_char
+                edits_made += 1
+        elif edit == "delete" and len(mangled) > 0:
+            mangled.pop(position)
+            edits_made += 1
+
+    return "".join(mangled)
+
+def get_random_edit_distance(
+    minimum: int, maximum: int, mean: float, dev: float
+) -> int:
+    sample = np.random.normal(loc=mean, scale=dev)
+    sample = int(sample)
+    return min(max(sample, minimum), maximum)
+
+def get_homologous_pair(
+    source: str, length: int
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    # Use more reasonable distance distribution
+    distance = get_random_edit_distance(1, min(length//4, 10), 3, 2)
+    mangled = mangle_string(source, distance)
+
+    # Verify actual distance and use it for training
+    actual_distance = ldistance(source, mangled)
+
+    return (
+        string_to_tensor(source, length),
+        string_to_tensor(mangled, length),
+        torch.tensor(float(actual_distance), dtype=torch.float),
+    )
+
+def get_non_homologous_pair(
+    length: int,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    source = random_str(length)
+    other = random_str(length)
+
+    # Ensure strings are actually different
+    max_attempts = 5
+    attempt = 0
+    while source == other and attempt < max_attempts:
+        other = random_str(length)
+        attempt += 1
+
+    distance = ldistance(source, other)
+
+    return (
+        string_to_tensor(source, length),
+        string_to_tensor(other, length),
+        torch.tensor(float(distance), dtype=torch.float),
+    )
+
+def squared_euclidean_distance(v1: torch.Tensor, v2: torch.Tensor) -> torch.Tensor:
+    return torch.sum((v1 - v2) ** 2, dim=1)
+
+def get_batch(
+    size: int, batch_size: int
+) -> list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
+    half_b = int(batch_size / 2)
+
+    # Generate diverse source strings for homologous pairs
+    h_pairs = []
+    for _ in range(half_b):
+        source = random_str(size)
+        h_pairs.append(get_homologous_pair(source, size))
+
+    non_h_pairs = [get_non_homologous_pair(size) for _ in range(half_b)]
+
+    # Shuffle the batch to prevent learning batch order patterns
+    all_pairs = h_pairs + non_h_pairs
+    np.random.shuffle(all_pairs)
+
+    return all_pairs
+
+def estimate_M(length: int, num_samples: int = 1000) -> float:
+    """Estimates M, the average Levenshtein distance for non-homologous pairs."""
+    total_distance = 0.0
+    for _ in range(num_samples):
+        _, _, dist_tensor = get_non_homologous_pair(length)
+        total_distance += dist_tensor.item()
+    return total_distance / num_samples
+
+def get_distances(
+    batch: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
+    model: torch.nn.Module,
+    M: float | None = None,
+    embedding_dim: int | None = None
+):
+    first: torch.Tensor = torch.stack([b[0] for b in batch])
+    first = model(first)
+
+    second: torch.Tensor = torch.stack([b[1] for b in batch])
+    second = model(second)
+
+    ds = torch.stack([b[2] for b in batch])
+
+    d_hats = squared_euclidean_distance(first, second)
+
+    if M is not None and embedding_dim is not None:
+        # r(n) = sqrt(M / (2n)) from paper Eq. 6
+        # We need r(n)^2 * d_hats, so (M / (2n)) * d_hats
+        scaling_factor_squared = M / (2 * embedding_dim)
+        d_hats = d_hats * scaling_factor_squared
+
+    return (d_hats, ds)
+
+def approximation_error(d_hat: torch.Tensor, d: torch.Tensor):
+    return torch.mean(torch.abs(d - d_hat))
+
+def get_loss(d_hat: torch.Tensor, d: torch.Tensor) -> torch.Tensor:
+    """
+    Wei et al. Poisson regression loss function
+    """
+    # Wei et al. Poisson regression with improved numerical stability
+    # PNLL(d̂, d) = d̂ - d * ln(d̂) with better handling of edge cases
+    epsilon = 1e-8
+    d_hat_stable = torch.clamp(d_hat, min=epsilon)
+    return torch.mean(d_hat_stable - d * torch.log(d_hat_stable))
+
+def validate_training_data(batch: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]]) -> dict:
+    """Validate and analyze training batch quality"""
+    distances = [b[2].item() for b in batch]
+
+    stats = {
+        'min_distance': min(distances),
+        'max_distance': max(distances),
+        'mean_distance': np.mean(distances),
+        'std_distance': np.std(distances),
+        'zero_distance_count': sum(1 for d in distances if d == 0),
+        'high_distance_count': sum(1 for d in distances if d > 15)
+    }
+
+    return stats
+
+def run_experiment(
+    embedding_dim: int,
+    model: torch.nn.Module,
+    learning_rate: float,
+    num_steps: int,
+    size: int,
+    batch_size: int,
+    use_gradient_clipping: bool = True,
+    max_grad_norm: float = 1.0,
+    distance_metric: str = "euclidean"
+):
+    """
+    Runs a training experiment with the given parameters and improved loss functions.
+    """
+    optimizer = AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-5)
+    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=200, gamma=0.5)
+    final_loss = 0.0
+    final_approx_error = 0.0
+
+    # Estimate M once at the beginning of the experiment
+    M_estimate = estimate_M(size)
+    print(f"Estimated M (average non-homologous distance): {M_estimate:.2f}")
+
+    for x in range(num_steps):
+        batch = get_batch(size, batch_size)
+
+        distances = get_distances(batch, model, distance_metric, M=M_estimate, embedding_dim=embedding_dim)
+        loss = get_loss(distances[0], distances[1])
+
+        if x % 10 == 0:
+            print(
+                f"step: {x}, loss: {loss.item()}, approx_error: {approximation_error(distances[0], distances[1]).item()}"
+            )
+
+        loss.backward()
+        optimizer.step()
+        scheduler.step()
+
+        final_loss = loss.item()
+        final_approx_error = approximation_error(distances[0], distances[1]).item()
+
+    return final_loss, final_approx_error
+
+if __name__ == "__main__":
+    embedding_dim = 140
+
+    model = EditDistanceModel(embedding_dim=embedding_dim)
+
+    final_loss, final_approx_error = run_experiment(
+        embedding_dim=embedding_dim,
+        model=model,
+        learning_rate=0.000817,
+        num_steps=1000,
+        size=80,
+        batch_size=32,
+        use_gradient_clipping=True,
+        max_grad_norm=2.463,
+        distance_metric="euclidean",
+    )
+
+    print(f"Final loss: {final_loss:.4f}")
+    print(f"Final approximation error: {final_approx_error:.4f}")
+
+    # Save the trained model
+    model_path = "megashtein_trained_model.pth"
+    torch.save(model.state_dict(), model_path)
+    print(f"\n model saved to: {model_path}")