Don't use ONNX because ZeroGPU doesn't support it.

app.py

@@ -65,7 +65,8 @@ def verbalize(text):
     parts = split_to_parts(text)
     verbalized = ''
     for part in parts:
-        verbalized += verbalizer.generate_text(part)
+        if part.strip():
+            verbalized += verbalizer.generate_text(part) + ' '
     return verbalized
 
 description = f'''

requirements.txt

@@ -13,5 +13,4 @@ git+https://github.com/patriotyk/ipa-uk.git
 git+https://github.com/patriotyk/styletts2-inference@105aed29fa1a7698d08d920986890e9bbd03447c
 spaces
 numpy<2
-huggingface_hub
-onnxruntime
+accelerate>=0.26.0

verbalizer.py

@@ -1,107 +1,39 @@
-import onnxruntime
-import numpy as np
-from transformers import AutoTokenizer
-from huggingface_hub import hf_hub_download
 
-verbalizer_model_name = "skypro1111/mbart-large-50-verbalization"
-
-def cache_model_from_hf(repo_id, model_dir="./"):
-    """Download ONNX models from HuggingFace Hub."""
-    files = ["onnx/encoder_model.onnx", "onnx/decoder_model.onnx", "onnx/decoder_model.onnx_data"]
-    
-    for file in files:
-        hf_hub_download(
-            repo_id=repo_id,
-            filename=file,
-            local_dir=model_dir,
-        )
+from transformers import MBartForConditionalGeneration, AutoTokenizer
 
+verbalizer_model_name = "skypro1111/mbart-large-50-verbalization"
 
 
 class Verbalizer():
     def __init__(self, device):
-        cache_model_from_hf(verbalizer_model_name)
-
-        print("Loading tokenizer...")
-        self.tokenizer = AutoTokenizer.from_pretrained(verbalizer_model_name)
-        self.tokenizer.src_lang = "uk_UA"
-        self.tokenizer.tgt_lang = "uk_UA"
+        self.device = device
 
-        print("Creating ONNX sessions...")
-        self.encoder_session = self.create_onnx_session("onnx/encoder_model.onnx", device=='cuda')
-        self.decoder_session = self.create_onnx_session("onnx/decoder_model.onnx", device=='cuda')
-
-    
-    def create_onnx_session(self, model_path, use_gpu=True):
-        """Create an ONNX inference session."""
-        session_options = onnxruntime.SessionOptions()
-        session_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
-        session_options.enable_mem_pattern = True
-        session_options.enable_mem_reuse = True
-        session_options.intra_op_num_threads = 8
-        #session_options.log_severity_level = 1
-        
-        cuda_provider_options = {
-            'device_id': 0,
-            'arena_extend_strategy': 'kSameAsRequested',
-            'gpu_mem_limit': 0,  # 0 means no limit
-            'cudnn_conv_algo_search': 'DEFAULT',
-            'do_copy_in_default_stream': True,
-        }
-        
-        if use_gpu and 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
-            providers = [('CUDAExecutionProvider', cuda_provider_options)]
-        else:
-            providers = ['CPUExecutionProvider']
-        
-        session = onnxruntime.InferenceSession(
-            model_path,
-            providers=providers,
-            sess_options=session_options
+        self.model = MBartForConditionalGeneration.from_pretrained(verbalizer_model_name,
+            low_cpu_mem_usage=True,
+            device_map=device,
         )
-        
-        return session
+        self.model.eval()
     
+        self.tokenizer = AutoTokenizer.from_pretrained(verbalizer_model_name)
+        self.tokenizer.src_lang = "uk_XX"
+        self.tokenizer.tgt_lang = "uk_XX"
     
     def generate_text(self, text):
         """Generate text for a single input."""
         # Prepare input
-        inputs = self.tokenizer(text, return_tensors="np", padding=True, truncation=True, max_length=512)
-        input_ids = inputs["input_ids"].astype(np.int64)
-        attention_mask = inputs["attention_mask"].astype(np.int64)
-        
-        # Run encoder
-        encoder_outputs = self.encoder_session.run(
-            output_names=["last_hidden_state"],
-            input_feed={
-                "input_ids": input_ids,
-                "attention_mask": attention_mask,
-            }
-        )[0]
-        
-        # Initialize decoder input
-        decoder_input_ids = np.array([[self.tokenizer.pad_token_id]], dtype=np.int64)
-        
-        # Generate sequence
-        for _ in range(512):
-            # Run decoder
-            decoder_outputs = self.decoder_session.run(
-                output_names=["logits"],
-                input_feed={
-                    "input_ids": decoder_input_ids,
-                    "encoder_hidden_states": encoder_outputs,
-                    "encoder_attention_mask": attention_mask,
-                }
-            )[0]
-            
-            # Get next token
-            next_token = decoder_outputs[:, -1:].argmax(axis=-1)
-            decoder_input_ids = np.concatenate([decoder_input_ids, next_token], axis=-1)
-            
-            # Check if sequence is complete
-            if self.tokenizer.eos_token_id in decoder_input_ids[0]:
-                break
-        
-        # Decode sequence
-        output_text = self.tokenizer.decode(decoder_input_ids[0], skip_special_tokens=True)
-        return output_text
+        input_text = "<verbalization>:" + text
+
+        encoded_input = self.tokenizer(
+            input_text,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=1024,
+        ).to(self.device)
+        output_ids = self.model.generate(
+            **encoded_input, max_length=1024, num_beams=5, early_stopping=True
+        )
+        normalized_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True)
+
+
+        return normalized_text.strip()