Susurrus/test_voxtral.py at main · CrispStrobe/Susurrus · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
#!/usr/bin/env python3
# test_voxtral.py - Test Voxtral integration

import argparse
import logging
import os
import sys

logging.basicConfig(level=logging.INFO, format="%(message)s")


def test_voxtral_local(audio_file, language=None):
    """Test local Voxtral inference"""
    print("\n" + "=" * 60)
    print("Testing Voxtral Local Inference")
    print("=" * 60)

    try:
        from voxtral_local import VoxtralLocal

        voxtral = VoxtralLocal(device="auto")

        print(f"\nTranscribing: {audio_file}")
        if language:
            print(f"Language: {language}")

        segments = voxtral.transcribe(audio_file, language=language)

        print(f"\n✅ Transcription complete! Found {len(segments)} segments:\n")

        for i, segment in enumerate(segments, 1):
            print(f"Segment {i}:")
            print(f"  [{segment['start']:.2f}s -> {segment['end']:.2f}s]")
            print(f"  {segment['text']}\n")

        voxtral.unload_model()
        return True

    except ImportError as e:
        print(f"\n❌ Import error: {e}")
        print("\nPlease run: ./install_voxtral.sh")
        return False
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback

        traceback.print_exc()
        return False


def test_voxtral_api(audio_file, api_key=None, language=None):
    """Test Voxtral API inference"""
    print("\n" + "=" * 60)
    print("Testing Voxtral API Inference")
    print("=" * 60)

    try:
        from voxtral_api import VoxtralAPI

        if not api_key:
            api_key = os.environ.get("MISTRAL_API_KEY")

        if not api_key:
            print("\n❌ No API key provided")
            print("Set MISTRAL_API_KEY environment variable or pass --api-key")
            print("Get your key from: https://console.mistral.ai/")
            return False

        voxtral = VoxtralAPI(api_key=api_key)

        print(f"\nTranscribing: {audio_file}")
        if language:
            print(f"Language: {language}")

        segments = voxtral.transcribe(audio_file, language=language)

        print(f"\n✅ Transcription complete! Found {len(segments)} segments:\n")

        for i, segment in enumerate(segments, 1):
            print(f"Segment {i}:")
            print(f"  [{segment['start']:.2f}s -> {segment['end']:.2f}s]")
            print(f"  {segment['text']}\n")

        return True

    except ImportError as e:
        print(f"\n❌ Import error: {e}")
        return False
    except Exception as e:
        print(f"\n❌ Error: {e}")
        import traceback

        traceback.print_exc()
        return False


def main():
    parser = argparse.ArgumentParser(description="Test Voxtral integration")
    parser.add_argument("audio_file", help="Path to audio file")
    parser.add_argument(
        "--backend", choices=["local", "api", "both"], default="both", help="Backend to test"
    )
    parser.add_argument("--api-key", help="Mistral API key (for API test)")
    parser.add_argument("--language", help="Language code (en, fr, es, de, it, pt, pl, nl)")

    args = parser.parse_args()

    if not os.path.exists(args.audio_file):
        print(f"❌ Audio file not found: {args.audio_file}")
        sys.exit(1)

    success = True

    if args.backend in ["local", "both"]:
        if not test_voxtral_local(args.audio_file, args.language):
            success = False

    if args.backend in ["api", "both"]:
        if not test_voxtral_api(args.audio_file, args.api_key, args.language):
            success = False

    if success:
        print("\n" + "=" * 60)
        print("✅ All tests passed!")
        print("=" * 60)
    else:
        print("\n" + "=" * 60)
        print("❌ Some tests failed")
        print("=" * 60)
        sys.exit(1)


if __name__ == "__main__":
    main()