Uploaded model

GPTTTS.py

@@ -0,0 +1,71 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer, AutoConfig
+from encodec import EncodecModel
+from encodec.utils import convert_audio
+import torch
+import torchaudio
+import re
+
+class GPTTTS(PreTrainedModel):
+    def __init__(self, *model_args, **model_kwargs):
+        super().__init__(AutoConfig.from_pretrained("Ekgren/distilgpt2-finetuned-common-voice"), *model_args, **model_kwargs)
+        self.model = AutoModelForCausalLM.from_pretrained("Ekgren/distilgpt2-finetuned-common-voice")
+        self.encodec_model = EncodecModel.encodec_model_24khz()
+        self.encodec_model.set_target_bandwidth(1.5)
+        self.sample_rate = self.encodec_model.sample_rate
+    
+    def forward(self, input_ids):
+        #decoded = tokenizer.decode(tokens[0], skip_special_tokens=True)
+        #decoded = input_text
+        # Get all audio_token_
+        #pattern = r'audio_token_(\d+)'
+        #audio_tokens = re.findall(pattern, decoded)
+        #audio_tokens = [int(token) for token in audio_tokens]
+
+        tokens = self.model.generate(input_ids, do_sample=True, max_length=1024, temperature=1, top_k=50, top_p=0.95)[0]
+        # Get all tokens which are larger than 50257, and subtract 50257 from them
+        audio_tokens = [token - 50257 for token in tokens if token > 50257]
+
+        number_of_codebooks = 2
+        number_of_samples = len(audio_tokens) // number_of_codebooks
+        frame = torch.zeros(1, number_of_codebooks, number_of_samples, dtype=torch.long)
+        for sample in range(number_of_samples):
+            for codebook in range(number_of_codebooks):
+              frame[0, codebook, sample] = audio_tokens[sample * number_of_codebooks + codebook]
+    
+        frames = [(frame, None)]
+
+        with torch.no_grad():
+            wav = self.encodec_model.decode(frames)
+    
+        return wav[0, :, :]
+
+
+class GPTTTSTokenizer(PreTrainedTokenizer):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.tokenizer = AutoTokenizer.from_pretrained("anforsm/distilgpt2-finetuned-common-voice")
+        self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
+    
+    def tokenize(self, text, *args, **kwargs):
+        prompt = f"text: {text}\nsound:"
+        return self.tokenizer(prompt, return_tensors="pt")
+    
+    def _tokenize(self, *args, **kwargs):
+        return self.tokenize(*args, **kwargs)
+    
+    def convert_tokens_to_ids(self, tokens):
+        return tokens["input_ids"]
+    
+    def convert_ids_to_tokens(self, ids):
+        return self.tokenizer.decode(ids[0], skip_special_tokens=True)
+    
+    def _batch_encode_plus(self, *args, **kwargs):
+        return self.tokenize(*args, **kwargs)
+    
+    def _encode_plus(self, *args, **kwargs):
+        return self.tokenize(*args, **kwargs)
+    
+    
+    def save_vocabulary(self, *args, **kwargs):
+        return self.tokenizer.save_vocabulary(*args, **kwargs)
+