import torch
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline

device = "cuda:0" if torch.cuda.is_available() else "cpu"
torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32

model_id = "openai/whisper-large-v3"
model = AutoModelForSpeechSeq2Seq.from_pretrained(
    model_id, torch_dtype=torch_dtype, low_cpu_mem_usage=True, use_safetensors=True
)
model.to(device)

processor = AutoProcessor.from_pretrained(model_id)

pipe = pipeline(
    "automatic-speech-recognition",
    model=model,
    tokenizer=processor.tokenizer,
    feature_extractor=processor.feature_extractor,
    max_new_tokens=128,
    chunk_length_s=30,
    batch_size=16,
    return_timestamps=True,
    torch_dtype=torch_dtype,
    device=device,
)

def speech_to_text(filename, language=None):
    try:
        result = pipe(filename, generate_kwargs={"language": language if language else "english"})
        return result['text']
    except Exception as e:
        return f"Error: {str(e)}"

if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Transcribe audio file using Whisper model")
    parser.add_argument("filename", help="Path to the audio file")
    parser.add_argument("--language", help="Language of the audio (default: english)", default="english")
    
    args = parser.parse_args()
    
    transcription = speech_to_text(args.filename, args.language)
    print("Transcription:", transcription)

python main.py <path to audio file> --language italian

config.json: 100%|█████████████████████████████████████████████████████████████████| 1.27k/1.27k [00:00<?, ?B/s]
C:\Users\simon\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\huggingface_hub\file_download.py:147: UserWarning: `huggingface_hub` cache-system uses 
symlinks by default to efficiently store duplicated files but your machine does not support them in C:\Users\simon\.cache\huggingface\hub\models--openai--whisper-large-v3. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.
To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. 
In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development
  warnings.warn(message)
Traceback (most recent call last):
  File "D:\Programming\Personal\SbobAI\main.py", line 8, in <module>
    model = AutoModelForSpeechSeq2Seq.from_pretrained(
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\simon\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\transformers\models\auto\auto_factory.py", line 564, in from_pretrained        
    return model_class.from_pretrained(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\simon\AppData\Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\transformers\modeling_utils.py", line 3372, in from_pretrained
    raise ImportError(
ImportError: Using `low_cpu_mem_usage=True` or a `device_map` requires Accelerate: `pip install 'accelerate>=0.26.0'`

model.safetensors:  15%|███████▊                                            | 461M/3.09G [01:03<05:35, 7.83MB/s]

model.safetensors: 100%|███████████████████████████████████████████████████| 3.09G/3.09G [06:42<00:00, 7.67MB/s]
generation_config.json: 100%|██████████████████████████████████████████████| 3.90k/3.90k [00:00<00:00, 7.82MB/s]
preprocessor_config.json: 100%|████████████████████████████████████████████████████████| 340/340 [00:00<?, ?B/s]
tokenizer_config.json: 100%|█████████████████████████████████████████████████| 283k/283k [00:00<00:00, 1.09MB/s]
vocab.json: 100%|██████████████████████████████████████████████████████████| 1.04M/1.04M [00:00<00:00, 1.97MB/s]
tokenizer.json: 100%|██████████████████████████████████████████████████████| 2.48M/2.48M [00:00<00:00, 4.93MB/s]
merges.txt: 100%|████████████████████████████████████████████████████████████| 494k/494k [00:00<00:00, 1.55MB/s]
normalizer.json: 100%|█████████████████████████████████████████████████████| 52.7k/52.7k [00:00<00:00, 45.3MB/s]
added_tokens.json: 100%|███████████████████████████████████████████████████| 34.6k/34.6k [00:00<00:00, 1.04MB/s]
special_tokens_map.json: 100%|█████████████████████████████████████████████████████| 2.07k/2.07k [00:00<?, ?B/s]
Transcription: Error: ffmpeg was not found but is required to load audio files from filename

Local\Packages\PythonSoftwareFoundation.Python.3.11_qbz5n2kfra8p0\LocalCache\local-packages\Python311\site-packages\transformers\models\whisper\generation_whisper.py:496: FutureWarning: The input name `inputs` is deprecated. Please make sure to use `input_features` instead.
  warnings.warn(
You have passed language=italian, but also have set `forced_decoder_ids` to [[1, None], [2, 50360]] which creates a conflict. `forced_decoder_ids` will be ignored in favor of language=italian.
The attention mask is not set and cannot be inferred from input because pad token is same as eos token. As a consequence, you may observe unexpected behavior. Please pass your input's `attention_mask` to obtain reliable results.