""" 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()