Update yarngpt/generate.py

yarngpt/generate.py

@@ -9,6 +9,7 @@ from huggingface_hub import hf_hub_download
 import warnings
 import scipy.io.wavfile as wav
 from datetime import datetime
+import json
 
 # Configure logging
 logging.basicConfig(level=logging.INFO, 
@@ -16,7 +17,7 @@ logging.basicConfig(level=logging.INFO,
 logger = logging.getLogger(__name__)
 
 # Constants
-INIT_TIMESTAMP = "2025-05-21 02:08:00"
+INIT_TIMESTAMP = "2025-05-21 02:21:23"
 CURRENT_USER = "Abdulhameed556"
 
 class TextToSpeech:
@@ -26,40 +27,30 @@ class TextToSpeech:
         self.processor_name_or_path = processor_name_or_path or model_name_or_path
         self.init_time = INIT_TIMESTAMP
         self.user = CURRENT_USER
+        self.cache_dir = "/code/cache"
         
         logger.info(f"Initializing TextToSpeech with model: {model_name_or_path}")
         
         try:
+            # Create cache directory if it doesn't exist
+            os.makedirs(self.cache_dir, exist_ok=True)
+            
+            # Create tokenizer files locally if they don't exist
+            self._create_tokenizer_files()
+            
             # Initialize configuration
             config = Speech2Text2Config.from_pretrained(
                 pretrained_model_name_or_path=self.model_name_or_path,
-                cache_dir="/code/cache",
-                token=os.getenv('HF_TOKEN'),
-                trust_remote_code=True
+                cache_dir=self.cache_dir,
+                token=os.getenv('HF_TOKEN')
             )
             
-            # Download tokenizer files
-            logger.info("Downloading tokenizer files...")
-            tokenizer_files = ["tokenizer_config.json", "special_tokens_map.json", "vocab.json"]
-            for file in tokenizer_files:
-                try:
-                    hf_hub_download(
-                        repo_id=self.model_name_or_path,
-                        filename=file,
-                        cache_dir="/code/cache",
-                        token=os.getenv('HF_TOKEN')
-                    )
-                except Exception as e:
-                    logger.warning(f"Could not download {file}: {e}")
-            
             # Initialize tokenizer
             logger.info("Loading tokenizer...")
             self.tokenizer = AutoTokenizer.from_pretrained(
-                self.model_name_or_path,
+                self.cache_dir,  # Use local cache directory
                 config=config,
-                cache_dir="/code/cache",
-                token=os.getenv('HF_TOKEN'),
-                trust_remote_code=True
+                token=os.getenv('HF_TOKEN')
             )
             
             # Initialize model
@@ -70,26 +61,55 @@ class TextToSpeech:
             self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
                 self.model_name_or_path,
                 config=config,
-                cache_dir="/code/cache",
-                token=os.getenv('HF_TOKEN'),
-                trust_remote_code=True
+                cache_dir=self.cache_dir,
+                token=os.getenv('HF_TOKEN')
             ).to(self.device)
             
-            # Load processor
-            logger.info("Loading processor...")
-            self.processor = AutoProcessor.from_pretrained(
-                self.model_name_or_path,
-                cache_dir="/code/cache",
-                token=os.getenv('HF_TOKEN'),
-                trust_remote_code=True
-            )
-            
             logger.info("Model initialization complete")
             
         except Exception as e:
             logger.error(f"Error initializing TextToSpeech: {e}")
             raise
 
+    def _create_tokenizer_files(self):
+        """Create necessary tokenizer files in cache directory."""
+        tokenizer_files = {
+            "tokenizer_config.json": {
+                "name_or_path": self.model_name_or_path,
+                "padding_side": "right",
+                "truncation_side": "right",
+                "model_max_length": 1024,
+                "bos_token": "<s>",
+                "eos_token": "</s>",
+                "unk_token": "<unk>",
+                "pad_token": "<pad>",
+                "mask_token": "<mask>",
+                "special_tokens_map_file": "special_tokens_map.json",
+                "tokenizer_class": "Speech2Text2Tokenizer"
+            },
+            "special_tokens_map.json": {
+                "bos_token": "<s>",
+                "eos_token": "</s>",
+                "pad_token": "<pad>",
+                "unk_token": "<unk>",
+                "mask_token": "<mask>"
+            },
+            "vocab.json": {
+                "<s>": 0,
+                "<pad>": 1,
+                "</s>": 2,
+                "<unk>": 3,
+                "<mask>": 4
+            }
+        }
+        
+        logger.info("Creating tokenizer files in cache directory...")
+        for filename, content in tokenizer_files.items():
+            filepath = os.path.join(self.cache_dir, filename)
+            with open(filepath, 'w', encoding='utf-8') as f:
+                json.dump(content, f, indent=2)
+            logger.info(f"Created {filename}")
+
     def tts(self, text, speed=1.0):
         """Generate speech from text."""
         try:
@@ -114,7 +134,7 @@ class TextToSpeech:
                 )
             
             # Convert to audio
-            audio = self.processor.batch_decode(output, skip_special_tokens=True)[0]
+            audio = output[0].cpu().numpy()
             
             # Apply speed adjustment if needed
             if speed != 1.0: