model/upload_to_huggingface.py

"""
Upload Fine-tuned TinyLlama Model to Hugging Face Hub

This script uploads a fine-tuned TinyLlama model to the Hugging Face Hub.
It handles authentication, model card creation, and repository management.

Usage:
    python upload_to_huggingface.py --model_dir ./mt564_tinyllama_model --repo_name username/mt564-tinyllama
"""

import os
import argparse
import logging
from datetime import datetime
from typing import Optional, List, Dict, Any

logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

def parse_args():
    parser = argparse.ArgumentParser(description="Upload fine-tuned TinyLlama model to Hugging Face Hub")
    parser.add_argument(
        "--model_dir",
        type=str,
        required=True,
        help="Directory containing the fine-tuned model"
    )
    parser.add_argument(
        "--repo_name",
        type=str,
        required=True,
        help="Name for the Hugging Face repository (format: username/repo-name)"
    )
    parser.add_argument(
        "--commit_message",
        type=str,
        default=f"Upload fine-tuned TinyLlama model - {datetime.now().strftime('%Y-%m-%d')}",
        help="Commit message for the model upload"
    )
    parser.add_argument(
        "--private",
        action="store_true",
        help="Make the repository private"
    )
    parser.add_argument(
        "--create_model_card",
        action="store_true",
        default=True,
        help="Create a model card README.md"
    )
    parser.add_argument(
        "--base_model",
        type=str,
        default="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        help="Base model name used for fine-tuning"
    )
    parser.add_argument(
        "--tags",
        type=str,
        nargs="+",
        default=["swift", "mt564", "financial", "tinyllama", "finance"],
        help="Tags for the model"
    )
    return parser.parse_args()

def create_model_card(
    base_model: str,
    repo_name: str,
    dataset_info: str = "SWIFT MT564 documentation",
    tags: List[str] = None,
    training_details: Dict[str, Any] = None
) -> str:
    """
    Create a model card for the Hugging Face Hub
    
    Args:
        base_model: Name of the base model used for fine-tuning
        repo_name: Name of the Hugging Face repository
        dataset_info: Information about the dataset used
        tags: Tags for the model
        training_details: Dictionary with training hyperparameters
        
    Returns:
        Model card content as a string
    """
    if tags is None:
        tags = ["swift", "mt564", "finance", "tinyllama"]
    
    if training_details is None:
        training_details = {
            "epochs": 3,
            "learning_rate": "2e-5",
            "batch_size": 2,
            "gradient_accumulation_steps": 4,
            "training_date": datetime.now().strftime("%Y-%m-%d")
        }
    
    repo_owner, repo_id = repo_name.split('/')
    
    model_card = f"""---
language: en
license: apache-2.0
tags:
{chr(10).join([f'- {tag}' for tag in tags])}
datasets:
- custom
metrics:
- accuracy
---

# {repo_id}

This is a fine-tuned version of [{base_model}](https://huggingface.co/{base_model}) specialized for understanding SWIFT MT564 message formats and financial documentation.

## Model Description

This model was fine-tuned on SWIFT MT564 documentation to help financial professionals understand and work with Corporate Action Notification messages. It can answer questions about message structure, field specifications, and usage guidelines for MT564 messages.

### Base Model

- **Base Model**: {base_model}
- **Model Type**: TinyLlama
- **Language**: English
- **Fine-tuning Focus**: SWIFT financial messaging formats, particularly MT564

## Training Data

The model was fine-tuned on the following data:
- {dataset_info}
- The data includes message specifications, field descriptions, sequence structures, and usage guidelines

## Training Procedure

The model was fine-tuned with the following parameters:
- **Epochs**: {training_details['epochs']}
- **Learning Rate**: {training_details['learning_rate']}
- **Batch Size**: {training_details['batch_size']}
- **Gradient Accumulation Steps**: {training_details['gradient_accumulation_steps']}
- **Training Date**: {training_details['training_date']}

## Intended Use & Limitations

This model is specifically designed to:
- Answer questions about SWIFT MT564 message formats
- Assist with understanding Corporate Action Notifications
- Help parse and interpret MT564 messages

**Limitations**:
- This model specializes in MT564 and may have limited knowledge of other SWIFT message types
- The model should not be used for generating actual SWIFT messages for production systems
- Always verify critical financial information with official SWIFT documentation

## Usage

```python
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load model and tokenizer
model = AutoModelForCausalLM.from_pretrained("{repo_name}")
tokenizer = AutoTokenizer.from_pretrained("{repo_name}")

# Format prompt for the chat model
prompt = "<|im_start|>user\\nWhat is the purpose of Sequence A in MT564 messages?<|im_end|>\\n<|im_start|>assistant\\n"

# Tokenize and generate response
inputs = tokenizer(prompt, return_tensors="pt")
outputs = model.generate(inputs.input_ids, max_new_tokens=256, temperature=0.7)
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(response)
```

## Citation & Contact

If you use this model, please cite:
```
@misc{{{repo_id.replace('-', '_').lower()}}},
  author = {{{repo_owner}}},
  title = {{{repo_id} - A fine-tuned TinyLlama model for SWIFT MT564 documentation}},
  year = {{{datetime.now().year}}},
  publisher = {Hugging Face},
  journal = {Hugging Face Repository},
  howpublished = {{https://huggingface.co/{repo_name}}},
}
```

For questions or feedback, please reach out through the [Hugging Face community](https://discuss.huggingface.co/) or the GitHub repository linked to this project.
"""
    return model_card

def upload_to_hub(
    model_dir: str,
    repo_name: str,
    commit_message: str = "Upload fine-tuned model",
    private: bool = False,
    create_card: bool = True,
    base_model: str = "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
    tags: List[str] = None
) -> bool:
    """
    Upload model to Hugging Face Hub
    
    Args:
        model_dir: Directory containing the fine-tuned model
        repo_name: Name for the Hugging Face repository (username/repo-name)
        commit_message: Commit message for the upload
        private: Whether to make the repository private
        create_card: Whether to create a model card
        base_model: Base model used for fine-tuning
        tags: Tags for the model
        
    Returns:
        Success status (True if upload was successful)
    """
    try:
        # Import libraries here so the script doesn't fail if they're not installed
        from huggingface_hub import HfApi, create_repo
        
        # Check if HUGGING_FACE_TOKEN environment variable is set
        token = os.environ.get("HUGGING_FACE_TOKEN")
        if not token:
            logger.error("HUGGING_FACE_TOKEN environment variable is not set.")
            logger.error("Set it using: export HUGGING_FACE_TOKEN=your_token_here")
            return False
        
        api = HfApi(token=token)
        logger.info(f"Authenticated with Hugging Face Hub")
        
        # Create repository if it doesn't exist
        try:
            repo_url = create_repo(
                repo_id=repo_name,
                private=private,
                token=token,
                exist_ok=True
            )
            logger.info(f"Repository created/accessed: {repo_url}")
        except Exception as e:
            logger.error(f"Error creating repository: {e}")
            return False
        
        # Create and save model card if requested
        if create_card:
            logger.info("Creating model card")
            model_card_content = create_model_card(
                base_model=base_model,
                repo_name=repo_name,
                tags=tags
            )
            
            model_card_path = os.path.join(model_dir, "README.md")
            with open(model_card_path, "w", encoding="utf-8") as f:
                f.write(model_card_content)
            logger.info(f"Model card saved to {model_card_path}")
        
        # Upload model to Hub
        logger.info(f"Uploading model from {model_dir} to {repo_name}")
        api.upload_folder(
            folder_path=model_dir,
            repo_id=repo_name,
            commit_message=commit_message
        )
        
        logger.info(f"Model successfully uploaded to {repo_name}")
        logger.info(f"View your model at: https://huggingface.co/{repo_name}")
        return True
    
    except ImportError as e:
        logger.error(f"Required libraries not installed: {e}")
        logger.error("Please install huggingface_hub: pip install huggingface_hub")
        return False
    except Exception as e:
        logger.error(f"Error uploading model: {e}")
        return False

def main():
    args = parse_args()
    
    # Check if model directory exists
    if not os.path.isdir(args.model_dir):
        logger.error(f"Model directory {args.model_dir} does not exist")
        return
    
    # Check if model files exist
    model_files = [
        "pytorch_model.bin", "config.json", 
        "generation_config.json", "tokenizer_config.json",
        "tokenizer.json", "special_tokens_map.json"
    ]
    missing_files = [f for f in model_files if not os.path.exists(os.path.join(args.model_dir, f))]
    
    if missing_files:
        logger.warning(f"The following model files are missing: {', '.join(missing_files)}")
        logger.warning("The model might be in a different format or incomplete")
    
    # Upload to Hugging Face Hub
    logger.info(f"Uploading model from {args.model_dir} to {args.repo_name}")
    success = upload_to_hub(
        model_dir=args.model_dir,
        repo_name=args.repo_name,
        commit_message=args.commit_message,
        private=args.private,
        create_card=args.create_model_card,
        base_model=args.base_model,
        tags=args.tags
    )
    
    if success:
        logger.info(f"Model upload complete! Your model is now available at: https://huggingface.co/{args.repo_name}")
        logger.info("You can use it with Transformers library:")
        logger.info(f"from transformers import AutoModelForCausalLM, AutoTokenizer")
        logger.info(f"model = AutoModelForCausalLM.from_pretrained('{args.repo_name}')")
        logger.info(f"tokenizer = AutoTokenizer.from_pretrained('{args.repo_name}')")
    else:
        logger.error("Model upload failed. Please check the error messages above.")

if __name__ == "__main__":
    main()