维石
2024-04-19 dee1354d0d984df21d16a2eba5bacec31bfb0b4b
funasr/auto/auto_model.py
@@ -175,6 +175,8 @@
            kwargs["token_list"] = tokenizer.token_list if hasattr(tokenizer, "token_list") else None
            kwargs["token_list"] = tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else kwargs["token_list"]
            vocab_size = len(kwargs["token_list"]) if kwargs["token_list"] is not None else -1
            if vocab_size == -1 and hasattr(tokenizer, "get_vocab_size"):
                vocab_size = tokenizer.get_vocab_size()
        else:
            vocab_size = -1
        kwargs["tokenizer"] = tokenizer
@@ -211,6 +213,9 @@
            else:
                print(f"error, init_param does not exist!: {init_param}")
        
        # fp16
        if kwargs.get("fp16", False):
            model.to(torch.float16)
        return model, kwargs
    
    def __call__(self, *args, **cfg):
@@ -412,7 +417,7 @@
            return_raw_text = kwargs.get('return_raw_text', False)
            # step.3 compute punc model
            if self.punc_model is not None:
                if not len(result["text"]):
                if not len(result["text"].strip()):
                    if return_raw_text:
                        result['raw_text'] = ''
                else: