import torch
from transformers import PreTrainedModel, PretrainedConfig
from huggingface_hub import push_to_hub_pytorch
from torchtext.vocab import build_vocab_from_iterator, GloVe
from torchtext.data.utils import get_tokenizer
import torch.nn as nn
import torch.nn.functional as F

# Define the configuration class
class ContactSharingConfig(PretrainedConfig):
    model_type = "contact_sharing"
    
    def __init__(
        self,
        vocab_size=0,
        embed_dim=600,
        num_filters=600,
        filter_sizes=[3, 4, 5, 6, 7, 8, 9, 10],
        lstm_hidden_dim=768,
        output_dim=2,
        dropout=0.5,
        pad_idx=0,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.embed_dim = embed_dim
        self.num_filters = num_filters
        self.filter_sizes = filter_sizes
        self.lstm_hidden_dim = lstm_hidden_dim
        self.output_dim = output_dim
        self.dropout = dropout
        self.pad_idx = pad_idx

# Define the model class
class ContactSharingClassifier(PreTrainedModel):
    config_class = ContactSharingConfig
    
    def __init__(self, config):
        super().__init__(config)
        self.embedding = nn.Embedding(config.vocab_size, config.embed_dim, padding_idx=config.pad_idx)
        self.lstm = nn.LSTM(config.embed_dim, config.lstm_hidden_dim, bidirectional=True, batch_first=True)
        self.convs = nn.ModuleList([
            nn.Conv1d(in_channels=config.lstm_hidden_dim*2, out_channels=config.num_filters, kernel_size=fs)
            for fs in config.filter_sizes
        ])
        self.fc1 = nn.Linear(len(config.filter_sizes) * config.num_filters, len(config.filter_sizes) * config.num_filters // 2)
        self.fc2 = nn.Linear(len(config.filter_sizes) * config.num_filters // 2, config.output_dim)
        self.dropout = nn.Dropout(config.dropout)
        self.layer_norm = nn.LayerNorm(len(config.filter_sizes) * config.num_filters)

    def forward(self, text):
        embedded = self.embedding(text)
        lstm_out, _ = self.lstm(embedded)
        lstm_out = lstm_out.permute(0, 2, 1)
        conved = [F.relu(conv(lstm_out)) for conv in self.convs]
        pooled = [F.max_pool1d(conv, conv.shape[2]).squeeze(2) for conv in conved]
        cat = self.dropout(torch.cat(pooled, dim=1))
        cat = self.layer_norm(cat)
        x = F.relu(self.fc1(cat))
        x = self.dropout(x)
        return self.fc2(x)

# Load vocabulary
vocab = torch.load('vocab.pth')

# Create configuration
config = ContactSharingConfig(vocab_size=len(vocab), pad_idx=vocab["<pad>"])

# Create model
model = ContactSharingClassifier(config)

# Load trained weights
model.load_state_dict(torch.load('contact_sharing_epoch_1.pth', map_location='cpu'))

# Push to Hub
push_to_hub_pytorch(
    model,
    repo_name="contact-sharing-classifier",
    organization=None,  # Set this to your organization name if you're pushing to an organization
    use_temp_dir=True
)

print("Model uploaded successfully to Hugging Face Hub!")