import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

class FoodProjectionHead(nn.Module):
    """
    Contrastive Food Projection Head for BAAI/bge-m3 (1024-dim).
    Projects general multi-lingual embeddings into a specialized culinary semantic space.
    """
    def __init__(self, dim: int = 1024, residual_scale: float = 0.5):
        super().__init__()
        self.residual_scale = residual_scale
        self.net = nn.Sequential(
            nn.Linear(dim, dim),
            nn.LayerNorm(dim),
            nn.GELU(),
            nn.Linear(dim, dim)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        res = x
        out = self.net(x)
        return F.normalize(res + self.residual_scale * out, p=2, dim=-1)

    @classmethod
    def from_pretrained(cls, repo_id: str = "ozertuu/bge-m3-food-projection-head", token: str = None):
        model = cls(dim=1024)
        try:
            weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", token=token)
            state_dict = load_file(weights_path)
        except Exception:
            weights_path = hf_hub_download(repo_id=repo_id, filename="food_projection_head.pt", token=token)
            state_dict = torch.load(weights_path, map_location="cpu")
        model.load_state_dict(state_dict)
        model.eval()
        return model