import torch
import os

def save_model_for_demo():
    """
    This script should be run after training to save the model and vocabularies
    in a format that can be loaded by the Hugging Face demo
    """
    
    # Assuming you have these from your training
    # word_vocab, tag_vocab, model
    
    # Save vocabularies
    vocabs = {
        'word_vocab': word_vocab,
        'tag_vocab': tag_vocab
    }
    torch.save(vocabs, 'vocab.pth')
    
    # Save model
    torch.save({
        'model_state_dict': model.state_dict(),
        'word_vocab_size': len(word_vocab),
        'tag_vocab_size': len(tag_vocab)
    }, 'model_GRU_fold_1.pth')
    
    print("Model and vocabularies saved successfully!")

if __name__ == "__main__":
    # You'll need to modify this to use your actual trained model and vocabs
    print("Please run this after training with your actual model and vocabularies")