# summarizer_tool.py # Import necessary classes from the transformers library # pipeline is a high-level API for using models for various tasks # AutoTokenizer and AutoModelForSeq2SeqLM are used to load specific pre-trained models import os import torch # Import torch to ensure it's available as a backend from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSpeechSeq2Seq, AutoModelForImageClassification, AutoModelForObjectDetection from PIL import Image from pydub import AudioSegment import soundfile as sf import numpy as np # For numerical operations with audio/image data import io # To handle in-memory audio conversions from torch.utils.data import DataLoader from datasets import load_dataset, Audio, concatenate_datasets, Dataset, DatasetDict from torch.optim import AdamW # <--- Corrected import for PyTorch's AdamW from accelerate import Accelerator from tqdm.auto import tqdm from transformers import ( AutoFeatureExtractor, AutoModelForSpeechSeq2Seq, Seq2SeqTrainingArguments, Seq2SeqTrainer, DataCollatorForSeq2Seq, AutoTokenizer, AutoModelForSequenceClassification, DataCollatorWithPadding ) import evaluate import logging import re from langchain_community.vectorstores import FAISS from langchain_community.embeddings import HuggingFaceEmbeddings from langchain.text_splitter import RecursiveCharacterTextSplitter from langchain_community.document_loaders import PyPDFLoader from langchain.chains import RetrievalQA from gtts import gTTS import tempfile import time # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # --- Configuration --- MODEL_NAME = "bert-base-uncased" DATASET_NAME = "glue" # Example: GLUE benchmark, you can choose any other SUBSET_NAME = "sst2" # Example: SST-2 (Stanford Sentiment Treebank) NUM_EPOCHS = 3 BATCH_SIZE = 8 LEARNING_RATE = 2e-5 MAX_SEQ_LENGTH = 128 SAVE_MODEL_EVERY_EPOCH = False # Set to True to save after each epoch # --- Configuration for text classification training --- MODEL_CHECKPOINT = "bert-base-uncased" # Or "distilbert-base-uncased" for a smaller model TRAIN_SPLIT = "train" NUM_TRAIN_EPOCHS = 3 # Number of epochs for training # --- 1. Load Dataset in Streaming Mode --- def load_dataset_streaming(dataset_name, subset_name=None, split="train"): logging.info(f"Loading dataset '{dataset_name}' (subset: {subset_name}) in streaming mode for split '{split}'...") try: if subset_name: dataset = load_dataset(dataset_name, subset_name, split=split, streaming=True) else: dataset = load_dataset(dataset_name, split=split, streaming=True) logging.info(f"Dataset '{dataset_name}' loaded successfully in streaming mode.") return dataset except Exception as e: logging.error(f"Error loading dataset '{dataset_name}': {e}") return None # --- 2. Data Preprocessing Function --- def preprocess_function(examples, tokenizer, max_seq_length): """Tokenizes and prepares text data for the model.""" return tokenizer(examples["sentence"], truncation=True, padding="max_length", max_length=max_seq_length) # --- 3. Main Training Loop --- def train_model_with_streaming_data(): accelerator = Accelerator() logging.info(f"Accelerator device: {accelerator.device}") logging.info(f"Loading dataset 'glue' (subset: sst2) in streaming mode for split '{TRAIN_SPLIT}'...") train_dataset = load_dataset('glue', 'sst2', split=TRAIN_SPLIT, streaming=True) # <-- This line defines 'train_dataset' logging.info("Dataset 'glue' loaded successfully in streaming mode.") # ... inside train_model_with_streaming_data function ... tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT) model = AutoModelForSequenceClassification.from_pretrained(MODEL_CHECKPOINT, num_labels=2) # MAKE SURE THIS 'process_function' DEFINITION IS PRESENT AND CORRECTLY INDENTED HERE: def process_function(examples): tokenized_examples = tokenizer(examples["sentence"], truncation=True) # Ensure the 'label' from the dataset is mapped to 'labels' for the model tokenized_examples["labels"] = examples["label"] return tokenized_examples # END OF process_function DEFINITION # Your code should then proceed to call .map() like this: tokenized_train_dataset = train_dataset.map( process_function, batched=True, remove_columns=["sentence", "idx", "label"] ) # Define the data collator specific for text classification here # This ensures 'labels' are correctly padded and converted to tensors for BertForSequenceClassification data_collator_for_classification = DataCollatorWithPadding(tokenizer=tokenizer) ## <-- ADD THIS LINE # Inside train_model_with_streaming_data function, before line 92 logging.info(f"Loading dataset 'glue' (subset: sst2) in streaming mode for split '{TRAIN_SPLIT}'...") train_dataset = load_dataset('glue', 'sst2', split=TRAIN_SPLIT, streaming=True) # <-- This line defines 'train_dataset' logging.info("Dataset 'glue' loaded successfully in streaming mode.") # Apply the processing function to the dataset tokenized_train_dataset = train_dataset.map( process_function, batched=True, remove_columns=["sentence", "idx", "label"] ) # Ensure this line is tokenized_train_dataset.with_format("torch") if you had the 'columns' TypeError fix tokenized_train_dataset = tokenized_train_dataset.with_format("torch") # DataLoader for streaming dataset train_dataloader = DataLoader( tokenized_train_dataset, batch_size=BATCH_SIZE, collate_fn=data_collator_for_classification ## <-- CHANGE THIS LINE to use the new collator ) optimizer = AdamW(model.parameters(), lr=LEARNING_RATE) accelerator = Accelerator() model, optimizer, train_dataloader = accelerator.prepare( model, optimizer, train_dataloader ) model.train() for epoch in range(NUM_TRAIN_EPOCHS): logging.info(f"Starting Epoch {epoch + 1}/{NUM_TRAIN_EPOCHS}") for i, batch in enumerate(tqdm(train_dataloader)): if "label" in batch: # This fix was for the previous TypeError, keep it. batch.pop("label") outputs = model(**batch) loss = outputs.loss accelerator.backward(loss) optimizer.step() optimizer.zero_grad() if i % 100 == 0: logging.info(f"Step {i}, Loss: {loss.item()}") # Load dataset (train and validation splits for full training loop) # Note: For streaming, you typically iterate over a single split. # If you want to use a validation set, you might load it separately or use a smaller, non-streaming eval set. # For simplicity, we'll stream only the 'train' split for this example. # For real continuous learning, you might have a system that appends new data to your "training stream". train_dataset = load_dataset_streaming(DATASET_NAME, SUBSET_NAME, split="train") if train_dataset is None: return # Preprocess the streaming dataset # We use a wrapper function for map with streaming to ensure it works correctly # with the IterableDataset and can handle features. # Note: map on IterableDataset works differently than on Dataset. # It applies transformations on the fly. tokenized_train_dataset = train_dataset.map( lambda examples: preprocess_function(examples, tokenizer, MAX_SEQ_LENGTH), batched=True, remove_columns=["sentence", "idx"] # Remove original columns if not needed by the model ) # Set the format for PyTorch tokenized_train_dataset = tokenized_train_dataset.with_format("torch") # Create DataLoader - remember IterableDataset has no len(), so DataLoader needs careful handling # For training, you'll want to shuffle. buffer_size is important for streaming shuffle. train_dataloader = DataLoader(tokenized_train_dataset.shuffle(seed=42, buffer_size=10_000), batch_size=BATCH_SIZE) optimizer = AdamW(model.parameters(), lr=LEARNING_RATE) # Prepare for distributed training with Accelerator model, optimizer, train_dataloader = accelerator.prepare( model, optimizer, train_dataloader ) # --- Training Loop --- model.train() total_steps = 0 # To track total steps across epochs for epoch in range(NUM_EPOCHS): logging.info(f"Starting Epoch {epoch + 1}/{NUM_EPOCHS}") # When streaming, set_epoch is crucial for reshuffling and deterministic behavior # (if you use a seed) if hasattr(train_dataloader.dataset, 'set_epoch'): train_dataloader.dataset.set_epoch(epoch) progress_bar = tqdm(train_dataloader, desc=f"Epoch {epoch + 1}", disable=not accelerator.is_local_main_process) for batch_idx, batch in enumerate(progress_bar): # Move batch to device (handled by accelerator.prepare, but good to be aware) # batch = {k: v.to(accelerator.device) for k, v in batch.items()} if "label" in batch: batch.pop("label") outputs = model(**batch) loss = outputs.loss accelerator.backward(loss) optimizer.step() optimizer.zero_grad() progress_bar.set_postfix({'loss': loss.item()}) total_steps += 1 # Optional: Log loss more frequently than epoch end if total_steps % 100 == 0: accelerator.log({"train_loss": loss.item()}, step=total_steps) logging.info(f"Epoch {epoch + 1} finished. Avg Loss: {loss.item()}") # Reports last batch loss, not avg if SAVE_MODEL_EVERY_EPOCH: output_dir = f"./model_epoch_{epoch + 1}" accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) unwrapped_model.save_pretrained(output_dir, save_function=accelerator.save) tokenizer.save_pretrained(output_dir) logging.info(f"Model saved to {output_dir}") logging.info("Training complete.") # --- Saving the final model --- final_output_dir = "./final_model" accelerator.wait_for_everyone() unwrapped_model = accelerator.unwrap_model(model) unwrapped_model.save_pretrained(final_output_dir, save_function=accelerator.save) tokenizer.save_pretrained(final_output_dir) logging.info(f"Final model saved to {final_output_dir}") # You can then load this model for inference # from transformers import pipeline # classifier = pipeline("sentiment-analysis", model=final_output_dir, tokenizer=final_output_dir) # print(classifier("This movie was great!")) if __name__ == "__main__": train_model_with_streaming_data() # Configure logging logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # --- Global Cache for Pipelines --- # This prevents reloading the same model multiple times _pipeline_cache = {} def get_pipeline(task_name, model_name=None, **kwargs): """ Retrieves a Hugging Face pipeline, caching it for efficiency. """ cache_key = f"{task_name}-{model_name}-{hash(frozenset(kwargs.items()))}" if cache_key not in _pipeline_cache: logging.info(f"Loading pipeline for task '{task_name}' with model '{model_name}'...") if model_name: _pipeline_cache[cache_key] = pipeline(task_name, model=model_name, **kwargs) else: _pipeline_cache[cache_key] = pipeline(task_name, **kwargs) # Uses default model for task logging.info(f"Pipeline '{task_name}' loaded.") return _pipeline_cache[cache_key] class AIToolDispatcher: def __init__(self): logging.info("Initializing AI Tool Dispatcher...") # Define default models for various tasks. # You can change these to other models from Hugging Face Hub. self.default_models = { "text-classification": "distilbert-base-uncased-finetuned-sst-2-english", "sentiment-analysis": "distilbert-base-uncased-finetuned-sst-2-english", # Same as text-classification for sentiment "summarization": "sshleifer/distilbart-cnn-12-6", "text-generation": "gpt2", "translation_en_to_fr": "Helsinki-NLP/opus-mt-en-fr", "image-classification": "google/vit-base-patch16-224", "object-detection": "facebook/detr-resnet-50", "automatic-speech-recognition": "openai/whisper-tiny.en", # For English ASR # "question-answering": "distilbert-base-uncased-distilled-squad", # Example for QA # "fill-mask": "bert-base-uncased", # Example for fill-mask # "token-classification": "dbmdz/bert-large-cased-finetuned-conll03-eng", # Example for NER } logging.info("Dispatcher initialized.") def _get_task_pipeline(self, task: str, model_name: str = None): """Helper to get a cached pipeline for a specific task.""" final_model_name = model_name if model_name else self.default_models.get(task) if not final_model_name: raise ValueError(f"No default model specified for task '{task}'. Please provide `model_name`.") return get_pipeline(task, model_name=final_model_name) def process_text(self, text: str, task: str = "text-classification", **kwargs): """Processes text input for a given NLP task.""" if not isinstance(text, str): raise TypeError("Text input must be a string.") logging.info(f"Processing text for task: {task}") pipeline = self._get_task_pipeline(task) result = pipeline(text, **kwargs) return result def process_image(self, image_path: str, task: str = "image-classification", **kwargs): """Processes image file input for a given computer vision task.""" if not os.path.exists(image_path): raise FileNotFoundError(f"Image file not found: {image_path}") logging.info(f"Processing image for task: {task}") try: image = Image.open(image_path) except Exception as e: raise ValueError(f"Could not open image file: {e}") pipeline = self._get_task_pipeline(task) result = pipeline(image, **kwargs) return result def process_audio(self, audio_path: str, task: str = "automatic-speech-recognition", **kwargs): """Processes audio file input for a given audio task.""" if not os.path.exists(audio_path): raise FileNotFoundError(f"Audio file not found: {audio_path}") logging.info(f"Processing audio for task: {task}") # Whisper models expect audio in a specific format (16kHz, mono, float32) # We'll use pydub and soundfile to convert if necessary try: audio = AudioSegment.from_file(audio_path) # Convert to mono, 16kHz audio = audio.set_channels(1).set_frame_rate(16000) # Export to a format soundfile can read in memory buffer = io.BytesIO() audio.export(buffer, format="wav") buffer.seek(0) # Read with soundfile to get numpy array and sample rate array, sampling_rate = sf.read(buffer) # Ensure it's float32 if array.dtype != np.float32: array = array.astype(np.float32) except Exception as e: logging.error(f"Error processing audio file for conversion: {e}") raise ValueError(f"Could not prepare audio file: {e}") pipeline = self._get_task_pipeline(task) result = pipeline(array.tolist(), sampling_rate=sampling_rate, **kwargs) # Pass as list or array, and sample_rate return result def process_input(self, input_data: str, task: str = None, **kwargs): """ Main entry point for the AI tool. Tries to determine input type and dispatches to the appropriate processing function. Args: input_data (str): Can be raw text or a file path (for image/audio/video). task (str, optional): The specific AI task to perform (e.g., "summarization", "object-detection", "automatic-speech-recognition"). REQUIRED for non-text inputs. For text, it defaults to "text-classification". **kwargs: Additional arguments to pass to the specific pipeline. Returns: dict or list: The result from the AI model. """ if not isinstance(input_data, str): raise TypeError("Input data must be a string (raw text or file path).") # --- Modality Detection --- if os.path.exists(input_data): # It's a file path file_extension = input_data.split('.')[-1].lower() if file_extension in ['jpg', 'jpeg', 'png', 'gif', 'bmp', 'tiff']: if not task: task = "image-classification" if not task.startswith("image-"): logging.warning( f"Task '{task}' may not be suitable for image input. Defaulting to '{task}' anyway.") return self.process_image(input_data, task=task, **kwargs) elif file_extension in ['mp3', 'wav', 'ogg', 'flac', 'm4a']: if not task: task = "automatic-speech-recognition" if not task.startswith("audio-") and not task.startswith("automatic-speech-recognition"): logging.warning( f"Task '{task}' may not be suitable for audio input. Defaulting to '{task}' anyway.") return self.process_audio(input_data, task=task, **kwargs) elif file_extension in ['mp4', 'avi', 'mov', 'mkv']: # Video processing is complex and often involves frame extraction + image/audio processing # For a true "all-in-one" model, this would need dedicated video models or # a pipeline to extract frames and audio, then pass to image/audio models. # This example will skip full video processing for simplicity, but you can build upon it. logging.warning("Video processing is highly complex and not fully implemented in this example. " "You'd typically extract frames/audio and process them separately.") raise NotImplementedError("Full video processing is beyond this generalized example.") else: raise ValueError(f"Unsupported file type: .{file_extension}. Or specify task for this file.") else: # Assume it's raw text if not a file path if not task: task = "text-classification" # Default task for text if task not in self.default_models and not kwargs.get('model'): logging.warning( f"No default model for task '{task}'. Using default model for text-classification if available.") if task not in ["text-classification", "sentiment-analysis", "summarization", "text-generation", "translation_en_to_fr"]: raise ValueError( f"Unknown text task: '{task}'. Please choose from supported text tasks or provide a model_name.") return self.process_text(input_data, task=task, **kwargs) # --- Example Usage --- if __name__ == "__main__": dispatcher = AIToolDispatcher() # --- Text Examples --- print("\n--- Text Examples ---") text_input = "The new movie was absolutely fantastic! I highly recommend it." print(f"Input: '{text_input}'") print(f"Sentiment: {dispatcher.process_input(text_input, task='sentiment-analysis')}") text_to_summarize = ( "Artificial intelligence (AI) is intelligence—perceiving, synthesizing, and inferring information—demonstrated by machines, as opposed to intelligence displayed by animals or humans. Example tasks in which AI is used include speech recognition, computer vision, translation, and others. AI applications include advanced web search engines, recommendation systems, understanding human speech, self-driving cars, and competing at the highest level in strategic game systems.") print(f"\nInput: '{text_to_summarize}'") print(f"Summary: {dispatcher.process_input(text_to_summarize, task='summarization', max_length=50, min_length=10)}") text_to_generate = "In a galaxy far, far away, a brave knight" print(f"\nInput: '{text_to_generate}'") generated_text = dispatcher.process_input(text_to_generate, task='text-generation', max_new_tokens=50, num_return_sequences=1) print(f"Generated Text: {generated_text[0]['generated_text']}") # --- Image Examples --- print("\n--- Image Examples ---") # To run image examples, you need to have image files. # Let's create dummy image files for demonstration purposes if they don't exist. # In a real scenario, replace these with actual paths to your images. dummy_image_path_1 = "dummy_cat_image.jpg" dummy_image_path_2 = "dummy_building_image.png" if not os.path.exists(dummy_image_path_1): print(f"Creating a dummy image file at {dummy_image_path_1} for demonstration.") try: # Minimal image creation with PIL img = Image.new('RGB', (60, 30), color='red') img.save(dummy_image_path_1) # You would replace this with a real image path to get meaningful results # e.g., dummy_image_path_1 = "path/to/your/cat_picture.jpg" # NOTE: For actual object detection/image classification, you NEED real, diverse images. # A plain red image will yield uninteresting results. except ImportError: print( "Pillow not installed. Skipping dummy image creation. Please install Pillow and provide real image paths.") dummy_image_path_1 = None except Exception as e: print(f"Could not create dummy image: {e}. Skipping image examples.") dummy_image_path_1 = None if dummy_image_path_1 and os.path.exists(dummy_image_path_1): print(f"\nImage Input: {dummy_image_path_1}") try: print(f"Image Classification: {dispatcher.process_input(dummy_image_path_1, task='image-classification')}") except Exception as e: print(f"Error during image classification: {e}") # Object detection might require larger or more complex images # The dummy image is too simple for meaningful object detection. # You'll need real images for this. # print(f"\nImage Input for Object Detection: {dummy_image_path_1}") # print(f"Object Detection: {dispatcher.process_input(dummy_image_path_1, task='object-detection')}") # --- Audio Examples --- print("\n--- Audio Examples ---") # To run audio examples, you need to have audio files. # Let's create a dummy audio file for demonstration. # NOTE: Requires `pydub` and `soundfile`, and `ffmpeg` installed on your system. dummy_audio_path = "dummy_audio.wav" if not os.path.exists(dummy_audio_path): print( f"Creating a dummy audio file at {dummy_audio_path} for demonstration (requires pydub, soundfile, ffmpeg).") try: from pydub.generators import Sine sine_wave = Sine(440).to_audio_segment(duration=1000) # 1 second of 440 Hz tone sine_wave.export(dummy_audio_path, format="wav") # For actual ASR, you NEED real speech audio. # e.g., dummy_audio_path = "path/to/your/speech_recording.mp3" except ImportError: print( "pydub not installed. Skipping dummy audio creation. Please install pydub, soundfile, and ffmpeg (system-wide) and provide real audio paths.") dummy_audio_path = None except Exception as e: print(f"Could not create dummy audio: {e}. Skipping audio examples.") dummy_audio_path = None if dummy_audio_path and os.path.exists(dummy_audio_path): print(f"\nAudio Input: {dummy_audio_path}") try: # Whisper-tiny.en is an English-only model. # If your audio contains different language, use a multilingual Whisper model like "openai/whisper-tiny". transcription = dispatcher.process_input(dummy_audio_path, task='automatic-speech-recognition') print(f"Audio Transcription: {transcription['text']}") except NotImplementedError: print("Skipping audio due to unimplemented feature or missing dependencies.") except Exception as e: print(f"Error during audio transcription: {e}") # --- Clean up dummy files --- if os.path.exists(dummy_image_path_1) and "dummy_cat_image.jpg" in dummy_image_path_1: os.remove(dummy_image_path_1) if os.path.exists(dummy_audio_path) and "dummy_audio.wav" in dummy_audio_path: os.remove(dummy_audio_path) logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # --- Configuration --- MODEL_CHECKPOINT = "openai/whisper-tiny.en" # Using English-specific tiny model # For larger, multilingual models, use "openai/whisper-tiny", "openai/whisper-base", etc. LANGUAGE_ABBREVIATION = "en" # Language of the dataset DATASET_NAME = "common_voice" DATASET_CONFIG = "ja" # Example for Japanese, change to "en" for English Common Voice SPLIT_TRAIN = "train" SPLIT_VALID = "validation" # Or "test" BATCH_SIZE = 16 # Adjust based on your GPU VRAM GRADIENT_ACCUMULATION_STEPS = 1 # Increase for larger effective batch size if VRAM is limited LEARNING_RATE = 1e-5 NUM_TRAIN_EPOCHS = 3 # Start small, increase for better results OUTPUT_DIR = "./whisper-tiny-finetuned" MAX_INPUT_LENGTH_IN_S = 30.0 # Max duration of audio in seconds the model can handle # --- 1. Load Feature Extractor, Tokenizer, and Model --- def load_asr_components(model_checkpoint: str): logging.info(f"Loading feature extractor, tokenizer, and model for {model_checkpoint}...") feature_extractor = AutoFeatureExtractor.from_pretrained(model_checkpoint) tokenizer = AutoTokenizer.from_pretrained(model_checkpoint, language=LANGUAGE_ABBREVIATION, task="transcribe") model = AutoModelForSpeechSeq2Seq.from_pretrained(model_checkpoint) # Set pad token id for generation (important for Beam Search) model.config.forced_decoder_ids = None model.config.suppress_tokens = [] # Remove suppressed tokens if you want to generate anything # Ensure model's generation config matches tokenizer's language & task model.config.decoder_start_token_id = tokenizer.bos_token_id model.config.pad_token_id = tokenizer.pad_token_id logging.info("Components loaded.") return feature_extractor, tokenizer, model # --- 2. Load and Preprocess Dataset --- def prepare_dataset( dataset_name: str, dataset_config: str, split: str, feature_extractor, tokenizer, max_input_length_in_s: float ): logging.info(f"Loading dataset '{dataset_name}' with config '{dataset_config}' split '{split}'...") # Load dataset with 'audio' feature type for automatic resampling dataset = load_dataset(dataset_name, dataset_config, split=split, trust_remote_code=True) # Cast audio column to a specific sampling rate (usually 16kHz for ASR models) dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate)) # Filter out examples that are too long (optional, but good for performance) dataset = dataset.filter( lambda example: example["audio"]["array"].shape[0] < max_input_length_in_s * feature_extractor.sampling_rate, num_proc=os.cpu_count() # Use multiple processes for faster filtering ) logging.info(f"Dataset loaded and casted. Number of examples: {len(dataset)}") # Preprocessing function def prepare_example(example): # Load and resample audio (handled by .cast_column("audio", Audio(...))) audio = example["audio"] # Compute log-mel spectrograms from the audio waveform input_features = feature_extractor( audio["array"], sampling_rate=audio["sampling_rate"] ).input_features[0] # Tokenize target text (transcription) # Normalize text: lowercase, remove punctuation (adjust as needed for your data) sentence = example["sentence"] # Simple text normalization (customize heavily based on dataset and language) # Example for English Common Voice: remove common noise, punctuation etc. # This normalization MUST match your model's expected input for the target text. text_normalized = re.sub(r"[^\w\s]", "", sentence).lower() labels = tokenizer(text_normalized).input_ids return {"input_features": input_features, "labels": labels} # Apply preprocessing logging.info("Applying preprocessing to dataset...") # Use map with batched=False as audio processing is per example processed_dataset = dataset.map( prepare_example, remove_columns=dataset.column_names, # Remove original columns, keep only processed features and labels num_proc=os.cpu_count() if os.cpu_count() else 1 # Use multiple processes for faster preprocessing ) logging.info("Preprocessing complete.") return processed_dataset # --- 3. Data Collator --- # Pad input features and labels to max length within a batch class DataCollatorSpeechSeq2SeqWithPadding(DataCollatorForSeq2Seq): def __call__(self, features): input_features = [{"input_features": feature["input_features"]} for feature in features] batch = self.tokenizer.feature_extractor.pad(input_features, return_tensors="pt") label_features = [{"input_ids": feature["labels"]} for feature in features] labels_batch = self.tokenizer.pad(label_features, return_tensors="pt") # replace padding with -100 to ignore loss correctly labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100) batch["labels"] = labels return batch # --- 4. Define Evaluation Metric --- def compute_metrics(pred, tokenizer): metric = evaluate.load("wer") # Load Word Error Rate metric pred_ids = pred.predictions label_ids = pred.label_ids # Replace -100 in labels as we did in the DataCollator label_ids[label_ids == -100] = tokenizer.pad_token_id pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True) label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True) wer = 100 * metric.compute(predictions=pred_str, references=label_str) return {"wer": wer} # --- Main Fine-tuning Function --- def fine_tune_asr_model(): feature_extractor, tokenizer, model = load_asr_components(MODEL_CHECKPOINT) if None in [feature_extractor, tokenizer, model]: return # Load and preprocess datasets (using a small fraction for demonstration) # For full training, remove .select(range(...)) try: train_dataset = prepare_dataset( DATASET_NAME, DATASET_CONFIG, SPLIT_TRAIN, feature_extractor, tokenizer, MAX_INPUT_LENGTH_IN_S ).select(range(500)) # Take first 500 examples for quick demo eval_dataset = prepare_dataset( DATASET_NAME, DATASET_CONFIG, SPLIT_VALID, feature_extractor, tokenizer, MAX_INPUT_LENGTH_IN_S ).select(range(100)) # Take first 100 examples for quick demo except Exception as e: logging.error(f"Error preparing dataset: {e}") logging.info("Make sure the dataset name/config is correct and you have internet access.") logging.info("Also ensure the 'audio' and 'sentence' columns exist in your chosen dataset.") return # Initialize data collator data_collator = DataCollatorSpeechSeq2SeqWithPadding(tokenizer=tokenizer) # Initialize training arguments training_args = Seq2SeqTrainingArguments( output_dir=OUTPUT_DIR, per_device_train_batch_size=BATCH_SIZE, gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS, learning_rate=LEARNING_RATE, num_train_epochs=NUM_TRAIN_EPOCHS, predict_with_generate=True, # Essential for ASR to generate sequences generation_max_length=MAX_INPUT_LENGTH_IN_S, # Max length of generated transcription fp16=torch.cuda.is_available(), # Use mixed precision if GPU is available save_steps=500, # Save checkpoint every 500 steps eval_steps=500, # Evaluate every 500 steps logging_steps=25, # Log every 25 steps report_to=["tensorboard"], # Log to TensorBoard for visualization load_best_model_at_end=True, # Load best model based on eval_metric metric_for_best_model="wer", # Use WER for best model selection greater_is_better=False, # Lower WER is better push_to_hub=False, # Set to True to push model to Hugging Face Hub (requires login) ) # Initialize Trainer trainer = Seq2SeqTrainer( model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset, tokenizer=tokenizer, feature_extractor=feature_extractor, data_collator=data_collator, compute_metrics=lambda p: compute_metrics(p, tokenizer), ) # Train the model logging.info("Starting model training...") trainer.train() logging.info("Training complete. Saving final model.") # Save the final model trainer.save_model(OUTPUT_DIR) tokenizer.save_pretrained(OUTPUT_DIR) feature_extractor.save_pretrained(OUTPUT_DIR) logging.info(f"Fine-tuned ASR model saved to {OUTPUT_DIR}") logging.info("\n--- How to use the fine-tuned model ---") logging.info(f"from transformers import pipeline") logging.info( f"transcriber = pipeline('automatic-speech-recognition', model='{OUTPUT_DIR}', tokenizer='{OUTPUT_DIR}', feature_extractor='{OUTPUT_DIR}')") logging.info(f"print(transcriber('path/to/your/audio.wav'))") if __name__ == "__main__": # Ensure you have ffmpeg installed system-wide for audio processing! # (Refer to previous instructions if you haven't) fine_tune_asr_model() # --- Pipeline Cache --- _pipeline_cache = {} def get_pipeline(task, model_name=None, **kwargs): key = f"{task}-{model_name}-{hash(frozenset(kwargs.items()))}" if key not in _pipeline_cache: _pipeline_cache[key] = pipeline(task, model=model_name, **kwargs) if model_name else pipeline(task, **kwargs) return _pipeline_cache[key] # --- Default Models --- def_models = { "sentiment-analysis": "distilbert-base-uncased-finetuned-sst-2-english", "summarization": "sshleifer/distilbart-cnn-12-6", "text-generation": "gpt2", "translation_en_to_fr": "Helsinki-NLP/opus-mt-en-fr", "image-classification": "google/vit-base-patch16-224", "object-detection": "facebook/detr-resnet-50", "automatic-speech-recognition": "openai/whisper-tiny.en" } # --- RAG Components --- def build_rag_qa(pdf_path): loader = PyPDFLoader(pdf_path) docs = loader.load() splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) split_docs = splitter.split_documents(docs) embeddings = HuggingFaceEmbeddings() vectorstore = FAISS.from_documents(split_docs, embeddings) return RetrievalQA.from_chain_type(llm=pipeline("text-generation"), retriever=vectorstore.as_retriever()) # --- Main Dispatcher --- class AllInOneDispatcher: def __init__(self): self.memory = [] def _pipe(self, task, model=None): return get_pipeline(task, model_name=model or def_models.get(task)) def _is_file(self, path): return os.path.exists(path) def handle_text(self, text, task="sentiment-analysis", **kwargs): result = self._pipe(task)(text, **kwargs) self.memory.append((task, text, result)) return result def handle_image(self, path, task="image-classification", **kwargs): image = Image.open(path) result = self._pipe(task)(image, **kwargs) self.memory.append((task, path, result)) return result def handle_audio(self, path, task="automatic-speech-recognition", **kwargs): audio = AudioSegment.from_file(path).set_channels(1).set_frame_rate(16000) buf = io.BytesIO() audio.export(buf, format="wav") buf.seek(0) data, sr = sf.read(buf) if data.dtype != np.float32: data = data.astype(np.float32) result = self._pipe(task)(data.tolist(), sampling_rate=sr) self.memory.append((task, path, result)) return result def handle_video(self, path): import cv2 frames, audio_file = [], "extracted_audio.wav" cap = cv2.VideoCapture(path) while cap.isOpened(): ret, frame = cap.read() if not ret: break frames.append(Image.fromarray(frame)) if len(frames) >= 5: break cap.release() os.system(f"ffmpeg -i {path} -q:a 0 -map a {audio_file} -y") image_result = self.handle_image(frames[0], task="image-classification") audio_result = self.handle_audio(audio_file) os.remove(audio_file) return {"image": image_result, "audio": audio_result} def handle_pdf(self, path): qa = build_rag_qa(path) return qa.run("Summarize this document") def handle_tts(self, text, lang='en'): tts = gTTS(text=text, lang=lang) temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3").name tts.save(temp_path) return temp_path def process(self, input_data, task=None, **kwargs): if self._is_file(input_data): ext = input_data.split('.')[-1].lower() if ext in ['jpg', 'jpeg', 'png']: return self.handle_image(input_data, task, **kwargs) elif ext in ['mp3', 'wav']: return self.handle_audio(input_data, task, **kwargs) elif ext in ['mp4', 'mov']: return self.handle_video(input_data) elif ext in ['pdf']: return self.handle_pdf(input_data) else: raise ValueError(f"Unsupported file type: {ext}") else: if task == "tts": return self.handle_tts(input_data) return self.handle_text(input_data, task, **kwargs) # --- Example Usage --- if __name__ == "__main__": ai = AllInOneDispatcher() print("Text:", ai.process("The weather is great today!", task="sentiment-analysis")) print("Summarize:", ai.process("Artificial intelligence is a broad field...", task="summarization")) print("TTS path:", ai.process("This is a test speech.", task="tts"))