We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
torchaudio: torch.quantization -> torch.ao.quantization
Summary:
This changes the imports in the `torchaudio` to include the new import locations.

```
codemod -d pytorch/audio --extensions py 'torch.quantization' 'torch.ao.quantization'
```

Differential Revision: D31302450



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py`
Content:
```
1 #!/usr/bin/env python3
2 import argparse
3 import logging
4 import os
5 
6 import torch
7 import torchaudio
8 from torchaudio.models.wav2vec2.utils.import_huggingface import import_huggingface_model
9 from greedy_decoder import Decoder
10 
11 _LG = logging.getLogger(__name__)
12 
13 
14 def _parse_args():
15     parser = argparse.ArgumentParser(
16         description=__doc__,
17     )
18     parser.add_argument(
19         '--model',
20         required=True,
21         help='Path to the input pretrained weight file.'
22     )
23     parser.add_argument(
24         '--output-path',
25         help='Path to the directory, where the Torchscript-ed pipelines are saved.',
26     )
27     parser.add_argument(
28         '--test-file',
29         help='Path to a test audio file.',
30     )
31     parser.add_argument(
32         '--quantize',
33         action='store_true',
34         help='Quantize the model.',
35     )
36     parser.add_argument(
37         '--debug',
38         action='store_true',
39         help=(
40             'When enabled, individual components are separately tested '
41             'for the numerical compatibility and TorchScript compatibility.'
42         )
43     )
44     return parser.parse_args()
45 
46 
47 class Loader(torch.nn.Module):
48     def forward(self, audio_path: str) -> torch.Tensor:
49         waveform, sample_rate = torchaudio.load(audio_path)
50         if sample_rate != 16000:
51             waveform = torchaudio.functional.resample(waveform, float(sample_rate), 16000.)
52         return waveform
53 
54 
55 class Encoder(torch.nn.Module):
56     def __init__(self, encoder: torch.nn.Module):
57         super().__init__()
58         self.encoder = encoder
59 
60     def forward(self, waveform: torch.Tensor) -> torch.Tensor:
61         result, _ = self.encoder(waveform)
62         return result[0]
63 
64 
65 def _get_model(model_id):
66     from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
67     tokenizer = Wav2Vec2Processor.from_pretrained(model_id).tokenizer
68     labels = [k for k, v in sorted(tokenizer.get_vocab().items(), key=lambda kv: kv[1])]
69     original = Wav2Vec2ForCTC.from_pretrained(model_id)
70     model = import_huggingface_model(original)
71     return model.eval(), labels
72 
73 
74 def _get_decoder(labels):
75     return Decoder(labels)
76 
77 
78 def _main():
79     args = _parse_args()
80     _init_logging(args.debug)
81     _LG.info('Loading model: %s', args.model)
82     model, labels = _get_model(args.model)
83     _LG.info('Labels: %s', labels)
84     _LG.info('Building pipeline')
85     loader = Loader()
86     encoder = Encoder(model)
87     decoder = _get_decoder(labels)
88     _LG.info(encoder)
89 
90     if args.quantize:
91         _LG.info('Quantizing the model')
92         model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()
93         encoder = torch.quantization.quantize_dynamic(
94             encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)
95         _LG.info(encoder)
96 
97     # test
98     if args.test_file:
99         _LG.info('Testing with %s', args.test_file)
100         waveform = loader(args.test_file)
101         emission = encoder(waveform)
102         transcript = decoder(emission)
103         _LG.info(transcript)
104 
105     torch.jit.script(loader).save(os.path.join(args.output_path, 'loader.zip'))
106     torch.jit.script(encoder).save(os.path.join(args.output_path, 'encoder.zip'))
107     torch.jit.script(decoder).save(os.path.join(args.output_path, 'decoder.zip'))
108 
109 
110 def _init_logging(debug=False):
111     level = logging.DEBUG if debug else logging.INFO
112     format_ = (
113         '%(message)s' if not debug else
114         '%(asctime)s: %(levelname)7s: %(funcName)10s: %(message)s'
115     )
116     logging.basicConfig(level=level, format=format_)
117 
118 
119 if __name__ == '__main__':
120     _main()
121
```

Path: `examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py`
Content:
```
1 #!/usr/bin/evn python3
2 """Build Speech Recognition pipeline based on fairseq's wav2vec2.0 and dump it to TorchScript file.
3 
4 To use this script, you need `fairseq`.
5 """
6 import os
7 import argparse
8 import logging
9 
10 import torch
11 from torch.utils.mobile_optimizer import optimize_for_mobile
12 import torchaudio
13 from torchaudio.models.wav2vec2.utils.import_fairseq import import_fairseq_model
14 import fairseq
15 
16 from greedy_decoder import Decoder
17 
18 _LG = logging.getLogger(__name__)
19 
20 
21 def _parse_args():
22     parser = argparse.ArgumentParser(
23         description=__doc__,
24     )
25     parser.add_argument(
26         '--model-file',
27         required=True,
28         help='Path to the input pretrained weight file.'
29     )
30     parser.add_argument(
31         '--dict-dir',
32         help=(
33             'Path to the directory in which `dict.ltr.txt` file is found. '
34             'Required only when the model is finetuned.'
35         )
36     )
37     parser.add_argument(
38         '--output-path',
39         help='Path to the directory, where the TorchScript-ed pipelines are saved.',
40     )
41     parser.add_argument(
42         '--test-file',
43         help='Path to a test audio file.',
44     )
45     parser.add_argument(
46         '--debug',
47         action='store_true',
48         help=(
49             'When enabled, individual components are separately tested '
50             'for the numerical compatibility and TorchScript compatibility.'
51         )
52     )
53     parser.add_argument(
54         '--quantize',
55         action='store_true',
56         help='Apply quantization to model.'
57     )
58     parser.add_argument(
59         '--optimize-for-mobile',
60         action='store_true',
61         help='Apply optmization for mobile.'
62     )
63     return parser.parse_args()
64 
65 
66 class Loader(torch.nn.Module):
67     def forward(self, audio_path: str) -> torch.Tensor:
68         waveform, sample_rate = torchaudio.load(audio_path)
69         if sample_rate != 16000:
70             waveform = torchaudio.functional.resample(waveform, float(sample_rate), 16000.)
71         return waveform
72 
73 
74 class Encoder(torch.nn.Module):
75     def __init__(self, encoder: torch.nn.Module):
76         super().__init__()
77         self.encoder = encoder
78 
79     def forward(self, waveform: torch.Tensor) -> torch.Tensor:
80         result, _ = self.encoder(waveform)
81         return result[0]
82 
83 
84 def _get_decoder():
85     labels = [
86         "<s>",
87         "<pad>",
88         "</s>",
89         "<unk>",
90         "|",
91         "E",
92         "T",
93         "A",
94         "O",
95         "N",
96         "I",
97         "H",
98         "S",
99         "R",
100         "D",
101         "L",
102         "U",
103         "M",
104         "W",
105         "C",
106         "F",
107         "G",
108         "Y",
109         "P",
110         "B",
111         "V",
112         "K",
113         "'",
114         "X",
115         "J",
116         "Q",
117         "Z",
118     ]
119     return Decoder(labels)
120 
121 
122 def _load_fairseq_model(input_file, data_dir=None):
123     overrides = {}
124     if data_dir:
125         overrides['data'] = data_dir
126 
127     model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(
128         [input_file], arg_overrides=overrides
129     )
130     model = model[0]
131     return model
132 
133 
134 def _get_model(model_file, dict_dir):
135     original = _load_fairseq_model(model_file, dict_dir)
136     model = import_fairseq_model(original.w2v_encoder)
137     return model
138 
139 
140 def _main():
141     args = _parse_args()
142     _init_logging(args.debug)
143     loader = Loader()
144     model = _get_model(args.model_file, args.dict_dir).eval()
145     encoder = Encoder(model)
146     decoder = _get_decoder()
147     _LG.info(encoder)
148 
149     if args.quantize:
150         _LG.info('Quantizing the model')
151         model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()
152         encoder = torch.quantization.quantize_dynamic(
153             encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)
154         _LG.info(encoder)
155 
156     # test
157     if args.test_file:
158         _LG.info('Testing with %s', args.test_file)
159         waveform = loader(args.test_file)
160         emission = encoder(waveform)
161         transcript = decoder(emission)
162         _LG.info(transcript)
163 
164     torch.jit.script(loader).save(os.path.join(args.output_path, 'loader.zip'))
165     torch.jit.script(decoder).save(os.path.join(args.output_path, 'decoder.zip'))
166     scripted = torch.jit.script(encoder)
167     if args.optimize_for_mobile:
168         scripted = optimize_for_mobile(scripted)
169     scripted.save(os.path.join(args.output_path, 'encoder.zip'))
170 
171 
172 def _init_logging(debug=False):
173     level = logging.DEBUG if debug else logging.INFO
174     format_ = (
175         '%(message)s' if not debug else
176         '%(asctime)s: %(levelname)7s: %(funcName)10s: %(message)s'
177     )
178     logging.basicConfig(level=level, format=format_)
179 
180 
181 if __name__ == '__main__':
182     _main()
183
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py b/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py
--- a/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py
+++ b/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py
@@ -6,6 +6,7 @@
 import os
 import argparse
 import logging
+from typing import Tuple
 
 import torch
 from torch.utils.mobile_optimizer import optimize_for_mobile
@@ -15,6 +16,12 @@
 
 from greedy_decoder import Decoder
 
+TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(".")[:2])
+if TORCH_VERSION >= (1, 10):
+    import torch.ao.quantization as tq
+else:
+    import torch.quantization as tq
+
 _LG = logging.getLogger(__name__)
 
 
@@ -149,7 +156,7 @@
     if args.quantize:
         _LG.info('Quantizing the model')
         model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()
-        encoder = torch.quantization.quantize_dynamic(
+        encoder = tq.quantize_dynamic(
             encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)
         _LG.info(encoder)
 
diff --git a/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py b/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py
--- a/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py
+++ b/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py
@@ -2,12 +2,19 @@
 import argparse
 import logging
 import os
+from typing import Tuple
 
 import torch
 import torchaudio
 from torchaudio.models.wav2vec2.utils.import_huggingface import import_huggingface_model
 from greedy_decoder import Decoder
 
+TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(".")[:2])
+if TORCH_VERSION >= (1, 10):
+    import torch.ao.quantization as tq
+else:
+    import torch.quantization as tq
+
 _LG = logging.getLogger(__name__)
 
 
@@ -90,7 +97,7 @@
     if args.quantize:
         _LG.info('Quantizing the model')
         model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()
-        encoder = torch.quantization.quantize_dynamic(
+        encoder = tq.quantize_dynamic(
             encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)
         _LG.info(encoder)

{"golden_diff": "diff --git a/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py b/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py\n--- a/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py\n+++ b/examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py\n@@ -6,6 +6,7 @@\n import os\n import argparse\n import logging\n+from typing import Tuple\n \n import torch\n from torch.utils.mobile_optimizer import optimize_for_mobile\n@@ -15,6 +16,12 @@\n \n from greedy_decoder import Decoder\n \n+TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(\".\")[:2])\n+if TORCH_VERSION >= (1, 10):\n+    import torch.ao.quantization as tq\n+else:\n+    import torch.quantization as tq\n+\n _LG = logging.getLogger(__name__)\n \n \n@@ -149,7 +156,7 @@\n     if args.quantize:\n         _LG.info('Quantizing the model')\n         model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()\n-        encoder = torch.quantization.quantize_dynamic(\n+        encoder = tq.quantize_dynamic(\n             encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)\n         _LG.info(encoder)\n \ndiff --git a/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py b/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py\n--- a/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py\n+++ b/examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py\n@@ -2,12 +2,19 @@\n import argparse\n import logging\n import os\n+from typing import Tuple\n \n import torch\n import torchaudio\n from torchaudio.models.wav2vec2.utils.import_huggingface import import_huggingface_model\n from greedy_decoder import Decoder\n \n+TORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(\".\")[:2])\n+if TORCH_VERSION >= (1, 10):\n+    import torch.ao.quantization as tq\n+else:\n+    import torch.quantization as tq\n+\n _LG = logging.getLogger(__name__)\n \n \n@@ -90,7 +97,7 @@\n     if args.quantize:\n         _LG.info('Quantizing the model')\n         model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()\n-        encoder = torch.quantization.quantize_dynamic(\n+        encoder = tq.quantize_dynamic(\n             encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)\n         _LG.info(encoder)\n", "issue": "torchaudio: torch.quantization -> torch.ao.quantization\nSummary:\nThis changes the imports in the `torchaudio` to include the new import locations.\n\n```\ncodemod -d pytorch/audio --extensions py 'torch.quantization' 'torch.ao.quantization'\n```\n\nDifferential Revision: D31302450\n\n\n", "before_files": [{"content": "#!/usr/bin/env python3\nimport argparse\nimport logging\nimport os\n\nimport torch\nimport torchaudio\nfrom torchaudio.models.wav2vec2.utils.import_huggingface import import_huggingface_model\nfrom greedy_decoder import Decoder\n\n_LG = logging.getLogger(__name__)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(\n        description=__doc__,\n    )\n    parser.add_argument(\n        '--model',\n        required=True,\n        help='Path to the input pretrained weight file.'\n    )\n    parser.add_argument(\n        '--output-path',\n        help='Path to the directory, where the Torchscript-ed pipelines are saved.',\n    )\n    parser.add_argument(\n        '--test-file',\n        help='Path to a test audio file.',\n    )\n    parser.add_argument(\n        '--quantize',\n        action='store_true',\n        help='Quantize the model.',\n    )\n    parser.add_argument(\n        '--debug',\n        action='store_true',\n        help=(\n            'When enabled, individual components are separately tested '\n            'for the numerical compatibility and TorchScript compatibility.'\n        )\n    )\n    return parser.parse_args()\n\n\nclass Loader(torch.nn.Module):\n    def forward(self, audio_path: str) -> torch.Tensor:\n        waveform, sample_rate = torchaudio.load(audio_path)\n        if sample_rate != 16000:\n            waveform = torchaudio.functional.resample(waveform, float(sample_rate), 16000.)\n        return waveform\n\n\nclass Encoder(torch.nn.Module):\n    def __init__(self, encoder: torch.nn.Module):\n        super().__init__()\n        self.encoder = encoder\n\n    def forward(self, waveform: torch.Tensor) -> torch.Tensor:\n        result, _ = self.encoder(waveform)\n        return result[0]\n\n\ndef _get_model(model_id):\n    from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor\n    tokenizer = Wav2Vec2Processor.from_pretrained(model_id).tokenizer\n    labels = [k for k, v in sorted(tokenizer.get_vocab().items(), key=lambda kv: kv[1])]\n    original = Wav2Vec2ForCTC.from_pretrained(model_id)\n    model = import_huggingface_model(original)\n    return model.eval(), labels\n\n\ndef _get_decoder(labels):\n    return Decoder(labels)\n\n\ndef _main():\n    args = _parse_args()\n    _init_logging(args.debug)\n    _LG.info('Loading model: %s', args.model)\n    model, labels = _get_model(args.model)\n    _LG.info('Labels: %s', labels)\n    _LG.info('Building pipeline')\n    loader = Loader()\n    encoder = Encoder(model)\n    decoder = _get_decoder(labels)\n    _LG.info(encoder)\n\n    if args.quantize:\n        _LG.info('Quantizing the model')\n        model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()\n        encoder = torch.quantization.quantize_dynamic(\n            encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)\n        _LG.info(encoder)\n\n    # test\n    if args.test_file:\n        _LG.info('Testing with %s', args.test_file)\n        waveform = loader(args.test_file)\n        emission = encoder(waveform)\n        transcript = decoder(emission)\n        _LG.info(transcript)\n\n    torch.jit.script(loader).save(os.path.join(args.output_path, 'loader.zip'))\n    torch.jit.script(encoder).save(os.path.join(args.output_path, 'encoder.zip'))\n    torch.jit.script(decoder).save(os.path.join(args.output_path, 'decoder.zip'))\n\n\ndef _init_logging(debug=False):\n    level = logging.DEBUG if debug else logging.INFO\n    format_ = (\n        '%(message)s' if not debug else\n        '%(asctime)s: %(levelname)7s: %(funcName)10s: %(message)s'\n    )\n    logging.basicConfig(level=level, format=format_)\n\n\nif __name__ == '__main__':\n    _main()\n", "path": "examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py"}, {"content": "#!/usr/bin/evn python3\n\"\"\"Build Speech Recognition pipeline based on fairseq's wav2vec2.0 and dump it to TorchScript file.\n\nTo use this script, you need `fairseq`.\n\"\"\"\nimport os\nimport argparse\nimport logging\n\nimport torch\nfrom torch.utils.mobile_optimizer import optimize_for_mobile\nimport torchaudio\nfrom torchaudio.models.wav2vec2.utils.import_fairseq import import_fairseq_model\nimport fairseq\n\nfrom greedy_decoder import Decoder\n\n_LG = logging.getLogger(__name__)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(\n        description=__doc__,\n    )\n    parser.add_argument(\n        '--model-file',\n        required=True,\n        help='Path to the input pretrained weight file.'\n    )\n    parser.add_argument(\n        '--dict-dir',\n        help=(\n            'Path to the directory in which `dict.ltr.txt` file is found. '\n            'Required only when the model is finetuned.'\n        )\n    )\n    parser.add_argument(\n        '--output-path',\n        help='Path to the directory, where the TorchScript-ed pipelines are saved.',\n    )\n    parser.add_argument(\n        '--test-file',\n        help='Path to a test audio file.',\n    )\n    parser.add_argument(\n        '--debug',\n        action='store_true',\n        help=(\n            'When enabled, individual components are separately tested '\n            'for the numerical compatibility and TorchScript compatibility.'\n        )\n    )\n    parser.add_argument(\n        '--quantize',\n        action='store_true',\n        help='Apply quantization to model.'\n    )\n    parser.add_argument(\n        '--optimize-for-mobile',\n        action='store_true',\n        help='Apply optmization for mobile.'\n    )\n    return parser.parse_args()\n\n\nclass Loader(torch.nn.Module):\n    def forward(self, audio_path: str) -> torch.Tensor:\n        waveform, sample_rate = torchaudio.load(audio_path)\n        if sample_rate != 16000:\n            waveform = torchaudio.functional.resample(waveform, float(sample_rate), 16000.)\n        return waveform\n\n\nclass Encoder(torch.nn.Module):\n    def __init__(self, encoder: torch.nn.Module):\n        super().__init__()\n        self.encoder = encoder\n\n    def forward(self, waveform: torch.Tensor) -> torch.Tensor:\n        result, _ = self.encoder(waveform)\n        return result[0]\n\n\ndef _get_decoder():\n    labels = [\n        \"<s>\",\n        \"<pad>\",\n        \"</s>\",\n        \"<unk>\",\n        \"|\",\n        \"E\",\n        \"T\",\n        \"A\",\n        \"O\",\n        \"N\",\n        \"I\",\n        \"H\",\n        \"S\",\n        \"R\",\n        \"D\",\n        \"L\",\n        \"U\",\n        \"M\",\n        \"W\",\n        \"C\",\n        \"F\",\n        \"G\",\n        \"Y\",\n        \"P\",\n        \"B\",\n        \"V\",\n        \"K\",\n        \"'\",\n        \"X\",\n        \"J\",\n        \"Q\",\n        \"Z\",\n    ]\n    return Decoder(labels)\n\n\ndef _load_fairseq_model(input_file, data_dir=None):\n    overrides = {}\n    if data_dir:\n        overrides['data'] = data_dir\n\n    model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(\n        [input_file], arg_overrides=overrides\n    )\n    model = model[0]\n    return model\n\n\ndef _get_model(model_file, dict_dir):\n    original = _load_fairseq_model(model_file, dict_dir)\n    model = import_fairseq_model(original.w2v_encoder)\n    return model\n\n\ndef _main():\n    args = _parse_args()\n    _init_logging(args.debug)\n    loader = Loader()\n    model = _get_model(args.model_file, args.dict_dir).eval()\n    encoder = Encoder(model)\n    decoder = _get_decoder()\n    _LG.info(encoder)\n\n    if args.quantize:\n        _LG.info('Quantizing the model')\n        model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()\n        encoder = torch.quantization.quantize_dynamic(\n            encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)\n        _LG.info(encoder)\n\n    # test\n    if args.test_file:\n        _LG.info('Testing with %s', args.test_file)\n        waveform = loader(args.test_file)\n        emission = encoder(waveform)\n        transcript = decoder(emission)\n        _LG.info(transcript)\n\n    torch.jit.script(loader).save(os.path.join(args.output_path, 'loader.zip'))\n    torch.jit.script(decoder).save(os.path.join(args.output_path, 'decoder.zip'))\n    scripted = torch.jit.script(encoder)\n    if args.optimize_for_mobile:\n        scripted = optimize_for_mobile(scripted)\n    scripted.save(os.path.join(args.output_path, 'encoder.zip'))\n\n\ndef _init_logging(debug=False):\n    level = logging.DEBUG if debug else logging.INFO\n    format_ = (\n        '%(message)s' if not debug else\n        '%(asctime)s: %(levelname)7s: %(funcName)10s: %(message)s'\n    )\n    logging.basicConfig(level=level, format=format_)\n\n\nif __name__ == '__main__':\n    _main()\n", "path": "examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py"}], "after_files": [{"content": "#!/usr/bin/env python3\nimport argparse\nimport logging\nimport os\nfrom typing import Tuple\n\nimport torch\nimport torchaudio\nfrom torchaudio.models.wav2vec2.utils.import_huggingface import import_huggingface_model\nfrom greedy_decoder import Decoder\n\nTORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(\".\")[:2])\nif TORCH_VERSION >= (1, 10):\n    import torch.ao.quantization as tq\nelse:\n    import torch.quantization as tq\n\n_LG = logging.getLogger(__name__)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(\n        description=__doc__,\n    )\n    parser.add_argument(\n        '--model',\n        required=True,\n        help='Path to the input pretrained weight file.'\n    )\n    parser.add_argument(\n        '--output-path',\n        help='Path to the directory, where the Torchscript-ed pipelines are saved.',\n    )\n    parser.add_argument(\n        '--test-file',\n        help='Path to a test audio file.',\n    )\n    parser.add_argument(\n        '--quantize',\n        action='store_true',\n        help='Quantize the model.',\n    )\n    parser.add_argument(\n        '--debug',\n        action='store_true',\n        help=(\n            'When enabled, individual components are separately tested '\n            'for the numerical compatibility and TorchScript compatibility.'\n        )\n    )\n    return parser.parse_args()\n\n\nclass Loader(torch.nn.Module):\n    def forward(self, audio_path: str) -> torch.Tensor:\n        waveform, sample_rate = torchaudio.load(audio_path)\n        if sample_rate != 16000:\n            waveform = torchaudio.functional.resample(waveform, float(sample_rate), 16000.)\n        return waveform\n\n\nclass Encoder(torch.nn.Module):\n    def __init__(self, encoder: torch.nn.Module):\n        super().__init__()\n        self.encoder = encoder\n\n    def forward(self, waveform: torch.Tensor) -> torch.Tensor:\n        result, _ = self.encoder(waveform)\n        return result[0]\n\n\ndef _get_model(model_id):\n    from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor\n    tokenizer = Wav2Vec2Processor.from_pretrained(model_id).tokenizer\n    labels = [k for k, v in sorted(tokenizer.get_vocab().items(), key=lambda kv: kv[1])]\n    original = Wav2Vec2ForCTC.from_pretrained(model_id)\n    model = import_huggingface_model(original)\n    return model.eval(), labels\n\n\ndef _get_decoder(labels):\n    return Decoder(labels)\n\n\ndef _main():\n    args = _parse_args()\n    _init_logging(args.debug)\n    _LG.info('Loading model: %s', args.model)\n    model, labels = _get_model(args.model)\n    _LG.info('Labels: %s', labels)\n    _LG.info('Building pipeline')\n    loader = Loader()\n    encoder = Encoder(model)\n    decoder = _get_decoder(labels)\n    _LG.info(encoder)\n\n    if args.quantize:\n        _LG.info('Quantizing the model')\n        model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()\n        encoder = tq.quantize_dynamic(\n            encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)\n        _LG.info(encoder)\n\n    # test\n    if args.test_file:\n        _LG.info('Testing with %s', args.test_file)\n        waveform = loader(args.test_file)\n        emission = encoder(waveform)\n        transcript = decoder(emission)\n        _LG.info(transcript)\n\n    torch.jit.script(loader).save(os.path.join(args.output_path, 'loader.zip'))\n    torch.jit.script(encoder).save(os.path.join(args.output_path, 'encoder.zip'))\n    torch.jit.script(decoder).save(os.path.join(args.output_path, 'decoder.zip'))\n\n\ndef _init_logging(debug=False):\n    level = logging.DEBUG if debug else logging.INFO\n    format_ = (\n        '%(message)s' if not debug else\n        '%(asctime)s: %(levelname)7s: %(funcName)10s: %(message)s'\n    )\n    logging.basicConfig(level=level, format=format_)\n\n\nif __name__ == '__main__':\n    _main()\n", "path": "examples/libtorchaudio/speech_recognition/build_pipeline_from_huggingface_transformers.py"}, {"content": "#!/usr/bin/evn python3\n\"\"\"Build Speech Recognition pipeline based on fairseq's wav2vec2.0 and dump it to TorchScript file.\n\nTo use this script, you need `fairseq`.\n\"\"\"\nimport os\nimport argparse\nimport logging\nfrom typing import Tuple\n\nimport torch\nfrom torch.utils.mobile_optimizer import optimize_for_mobile\nimport torchaudio\nfrom torchaudio.models.wav2vec2.utils.import_fairseq import import_fairseq_model\nimport fairseq\n\nfrom greedy_decoder import Decoder\n\nTORCH_VERSION: Tuple[int, ...] = tuple(int(x) for x in torch.__version__.split(\".\")[:2])\nif TORCH_VERSION >= (1, 10):\n    import torch.ao.quantization as tq\nelse:\n    import torch.quantization as tq\n\n_LG = logging.getLogger(__name__)\n\n\ndef _parse_args():\n    parser = argparse.ArgumentParser(\n        description=__doc__,\n    )\n    parser.add_argument(\n        '--model-file',\n        required=True,\n        help='Path to the input pretrained weight file.'\n    )\n    parser.add_argument(\n        '--dict-dir',\n        help=(\n            'Path to the directory in which `dict.ltr.txt` file is found. '\n            'Required only when the model is finetuned.'\n        )\n    )\n    parser.add_argument(\n        '--output-path',\n        help='Path to the directory, where the TorchScript-ed pipelines are saved.',\n    )\n    parser.add_argument(\n        '--test-file',\n        help='Path to a test audio file.',\n    )\n    parser.add_argument(\n        '--debug',\n        action='store_true',\n        help=(\n            'When enabled, individual components are separately tested '\n            'for the numerical compatibility and TorchScript compatibility.'\n        )\n    )\n    parser.add_argument(\n        '--quantize',\n        action='store_true',\n        help='Apply quantization to model.'\n    )\n    parser.add_argument(\n        '--optimize-for-mobile',\n        action='store_true',\n        help='Apply optmization for mobile.'\n    )\n    return parser.parse_args()\n\n\nclass Loader(torch.nn.Module):\n    def forward(self, audio_path: str) -> torch.Tensor:\n        waveform, sample_rate = torchaudio.load(audio_path)\n        if sample_rate != 16000:\n            waveform = torchaudio.functional.resample(waveform, float(sample_rate), 16000.)\n        return waveform\n\n\nclass Encoder(torch.nn.Module):\n    def __init__(self, encoder: torch.nn.Module):\n        super().__init__()\n        self.encoder = encoder\n\n    def forward(self, waveform: torch.Tensor) -> torch.Tensor:\n        result, _ = self.encoder(waveform)\n        return result[0]\n\n\ndef _get_decoder():\n    labels = [\n        \"<s>\",\n        \"<pad>\",\n        \"</s>\",\n        \"<unk>\",\n        \"|\",\n        \"E\",\n        \"T\",\n        \"A\",\n        \"O\",\n        \"N\",\n        \"I\",\n        \"H\",\n        \"S\",\n        \"R\",\n        \"D\",\n        \"L\",\n        \"U\",\n        \"M\",\n        \"W\",\n        \"C\",\n        \"F\",\n        \"G\",\n        \"Y\",\n        \"P\",\n        \"B\",\n        \"V\",\n        \"K\",\n        \"'\",\n        \"X\",\n        \"J\",\n        \"Q\",\n        \"Z\",\n    ]\n    return Decoder(labels)\n\n\ndef _load_fairseq_model(input_file, data_dir=None):\n    overrides = {}\n    if data_dir:\n        overrides['data'] = data_dir\n\n    model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task(\n        [input_file], arg_overrides=overrides\n    )\n    model = model[0]\n    return model\n\n\ndef _get_model(model_file, dict_dir):\n    original = _load_fairseq_model(model_file, dict_dir)\n    model = import_fairseq_model(original.w2v_encoder)\n    return model\n\n\ndef _main():\n    args = _parse_args()\n    _init_logging(args.debug)\n    loader = Loader()\n    model = _get_model(args.model_file, args.dict_dir).eval()\n    encoder = Encoder(model)\n    decoder = _get_decoder()\n    _LG.info(encoder)\n\n    if args.quantize:\n        _LG.info('Quantizing the model')\n        model.encoder.transformer.pos_conv_embed.__prepare_scriptable__()\n        encoder = tq.quantize_dynamic(\n            encoder, qconfig_spec={torch.nn.Linear}, dtype=torch.qint8)\n        _LG.info(encoder)\n\n    # test\n    if args.test_file:\n        _LG.info('Testing with %s', args.test_file)\n        waveform = loader(args.test_file)\n        emission = encoder(waveform)\n        transcript = decoder(emission)\n        _LG.info(transcript)\n\n    torch.jit.script(loader).save(os.path.join(args.output_path, 'loader.zip'))\n    torch.jit.script(decoder).save(os.path.join(args.output_path, 'decoder.zip'))\n    scripted = torch.jit.script(encoder)\n    if args.optimize_for_mobile:\n        scripted = optimize_for_mobile(scripted)\n    scripted.save(os.path.join(args.output_path, 'encoder.zip'))\n\n\ndef _init_logging(debug=False):\n    level = logging.DEBUG if debug else logging.INFO\n    format_ = (\n        '%(message)s' if not debug else\n        '%(asctime)s: %(levelname)7s: %(funcName)10s: %(message)s'\n    )\n    logging.basicConfig(level=level, format=format_)\n\n\nif __name__ == '__main__':\n    _main()\n", "path": "examples/libtorchaudio/speech_recognition/build_pipeline_from_fairseq.py"}]}