swift_setup.py

"""
Setup Script for SWIFT MT564 Documentation Assistant

This script guides you through setting up the project, installing dependencies,
downloading the TinyLlama model, and uploading your fine-tuned model to Hugging Face.

Usage:
    python setup.py --mode [install|download|upload] --options [options]
"""

import os
import sys
import argparse
import subprocess
import logging

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

def parse_args():
    parser = argparse.ArgumentParser(description="Setup SWIFT MT564 Documentation Assistant")
    parser.add_argument(
        "--mode",
        type=str,
        choices=["install", "download", "upload", "guide"],
        default="guide",
        help="Mode of operation: install dependencies, download model, upload model, or show guide"
    )
    parser.add_argument(
        "--model_name",
        type=str,
        default="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        help="Model name on Hugging Face Hub (for download mode)"
    )
    parser.add_argument(
        "--output_dir",
        type=str,
        default="./data/models",
        help="Directory to save the model (for download mode)"
    )
    parser.add_argument(
        "--model_dir",
        type=str,
        default="./mt564_tinyllama_model",
        help="Directory containing fine-tuned model (for upload mode)"
    )
    parser.add_argument(
        "--repo_name",
        type=str,
        help="Repository name on Hugging Face Hub (for upload mode)"
    )
    return parser.parse_args()

def install_dependencies():
    """Install required Python packages"""
    try:
        logger.info("Installing dependencies...")
        
        if os.path.exists("dependencies.txt"):
            subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "dependencies.txt"])
            logger.info("Dependencies installed successfully")
        else:
            logger.error("dependencies.txt file not found")
            logger.info("You need to install the following packages:")
            logger.info("- torch, transformers, datasets, accelerate")
            logger.info("- huggingface_hub, requests, beautifulsoup4, trafilatura")
            logger.info("- flask, PyPDF2, pdf2image, pytesseract")
            logger.info("Run: pip install torch transformers datasets huggingface_hub accelerate flask")
        
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Error installing dependencies: {e}")
        return False

def download_model(model_name, output_dir):
    """Download TinyLlama model"""
    try:
        logger.info(f"Downloading model {model_name}...")
        
        if not os.path.exists("model/download_tinyllama.py"):
            logger.error("download_tinyllama.py script not found")
            return False
        
        subprocess.check_call([
            sys.executable, 
            "model/download_tinyllama.py",
            "--model_name", model_name,
            "--output_dir", output_dir
        ])
        
        logger.info(f"Model downloaded successfully to {output_dir}")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Error downloading model: {e}")
        return False

def upload_model(model_dir, repo_name):
    """Upload fine-tuned model to Hugging Face Hub"""
    try:
        if not repo_name:
            logger.error("Repository name is required for upload mode")
            return False
        
        if not os.path.exists("model/upload_to_huggingface.py"):
            logger.error("upload_to_huggingface.py script not found")
            return False
        
        # Check for authentication token
        if "HUGGING_FACE_TOKEN" not in os.environ:
            logger.error("HUGGING_FACE_TOKEN environment variable not set")
            logger.info("You need to set your Hugging Face API token:")
            logger.info("export HUGGING_FACE_TOKEN=your_token_here")
            return False
        
        logger.info(f"Uploading model from {model_dir} to {repo_name}...")
        
        subprocess.check_call([
            sys.executable,
            "model/upload_to_huggingface.py",
            "--model_dir", model_dir,
            "--repo_name", repo_name
        ])
        
        logger.info(f"Model uploaded successfully to {repo_name}")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"Error uploading model: {e}")
        return False

def show_guide():
    """Show guide for using the SWIFT MT564 Documentation Assistant"""
    guide = """
=================================================
SWIFT MT564 Documentation Assistant - Setup Guide
=================================================

This project helps you build an AI assistant specialized in SWIFT MT564 messages.
Below is a step-by-step guide to setting up and using this project:

1. INSTALLATION
---------------
Install dependencies:
  python setup.py --mode install

2. DATA COLLECTION
-----------------
Run the ISO20022 scraper:
  python scrapers/iso20022_scraper.py --output_dir ./data/raw

Process the scraped data:
  python scrapers/data_processor.py --input_dir ./data/raw --output_dir ./data/processed

3. MODEL PREPARATION
-------------------
Download TinyLlama model:
  python setup.py --mode download --model_name TinyLlama/TinyLlama-1.1B-Chat-v1.0 --output_dir ./data/models

4. TRAINING
----------
Train the model on MT564 data:
  python train_mt564_model.py --model_name TinyLlama/TinyLlama-1.1B-Chat-v1.0 --training_data ./data/processed/mt564_training_data.json --output_dir ./mt564_tinyllama_model

5. DEPLOYMENT
------------
Upload model to Hugging Face Hub:
  # First set your Hugging Face token
  export HUGGING_FACE_TOKEN=your_token_here
  
  # Then upload the model
  python setup.py --mode upload --model_dir ./mt564_tinyllama_model --repo_name your-username/mt564-tinyllama

6. USING THE MODEL ON HUGGING FACE
---------------------------------
Create a Hugging Face Space:
  - Go to huggingface.co
  - Click "New Space"
  - Choose a Gradio or Streamlit template
  - Link to your uploaded model
  - Use the following code in your app.py:

```python
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load your model from Hugging Face Hub
model_name = "your-username/mt564-tinyllama"
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

def answer_question(question):
    # Format prompt for the chat model
    prompt = f"<|im_start|>user\\n{question}<|im_end|>\\n<|im_start|>assistant\\n"
    
    # 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)
    
    # Extract assistant's response
    try:
        assistant_response = response.split("<|im_start|>assistant\\n")[1].split("<|im_end|>")[0]
    except:
        assistant_response = response
    
    return assistant_response

# Create Gradio interface
iface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(lines=3, placeholder="Ask a question about SWIFT MT564..."),
    outputs=gr.Textbox(),
    title="SWIFT MT564 Documentation Assistant",
    description="Ask questions about SWIFT MT564 message formats and specifications"
)

iface.launch()
```

For more details and advanced usage, refer to the README.md.
"""
    print(guide)
    return True

def main():
    args = parse_args()
    
    if args.mode == "install":
        success = install_dependencies()
    elif args.mode == "download":
        success = download_model(args.model_name, args.output_dir)
    elif args.mode == "upload":
        success = upload_model(args.model_dir, args.repo_name)
    else:  # guide
        success = show_guide()
    
    if success:
        logger.info(f"{args.mode.capitalize()} completed successfully")
    else:
        logger.error(f"{args.mode.capitalize()} failed")
        sys.exit(1)

if __name__ == "__main__":
    main()