Delete CHATTS/infer/api.py

CHATTS/infer/api.py

@@ -1,125 +0,0 @@
-
-import torch
-import torch.nn.functional as F
-from transformers.generation import TopKLogitsWarper, TopPLogitsWarper
-from ..utils.infer_utils import CustomRepetitionPenaltyLogitsProcessorRepeat
-
-def infer_code(
-    models,
-    text, 
-    spk_emb = None,
-    top_P = 0.7, 
-    top_K = 20, 
-    temperature = 0.3, 
-    repetition_penalty = 1.05,
-    max_new_token = 2048,
-    **kwargs
-):
-    
-    device = next(models['gpt'].parameters()).device
-    
-    if not isinstance(text, list): 
-        text = [text]
-        
-    if not isinstance(temperature, list):
-        temperature = [temperature] * models['gpt'].num_vq
-    
-    if spk_emb is not None:
-        text = [f'[Stts][spk_emb]{i}[uv_break][Ptts]' for i in text] 
-    else:
-        text = [f'[Stts][empty_spk]{i}[uv_break][Ptts]' for i in text]
-    
-    text_token = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True).to(device)
-    input_ids = text_token['input_ids'][...,None].expand(-1, -1, models['gpt'].num_vq)
-    text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device)
-    
-    inputs = {
-        'input_ids': input_ids,
-        'text_mask': text_mask,
-        'attention_mask': text_token['attention_mask'],
-    }
-
-    emb = models['gpt'].get_emb(**inputs)
-    if spk_emb is not None:
-        emb[inputs['input_ids'][..., 0] == models['tokenizer'].convert_tokens_to_ids('[spk_emb]')] = \
-            F.normalize(spk_emb.to(device).to(emb.dtype)[None].expand(len(text), -1), p=2.0, dim=1, eps=1e-12)  
-    
-    num_code = models['gpt'].emb_code[0].num_embeddings - 1
-    
-    LogitsWarpers = []
-    if top_P is not None:
-        LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
-    if top_K is not None:
-        LogitsWarpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
-        
-    LogitsProcessors = []
-    if repetition_penalty is not None and repetition_penalty != 1:
-        LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(\
-            repetition_penalty, num_code, 16))
-    
-    result = models['gpt'].generate(
-        emb, inputs['input_ids'], 
-        temperature = torch.tensor(temperature, device=device), 
-        attention_mask = inputs['attention_mask'],
-        LogitsWarpers = LogitsWarpers,
-        LogitsProcessors = LogitsProcessors,
-        eos_token = num_code, 
-        max_new_token = max_new_token, 
-        infer_text = False,
-        **kwargs
-    )
-    
-    return result
-
-
-def refine_text(
-    models, 
-    text,
-    top_P = 0.7, 
-    top_K = 20, 
-    temperature = 0.7, 
-    repetition_penalty = 1.0,
-    max_new_token = 384,
-    prompt = '',
-    **kwargs
-):
-    
-    device = next(models['gpt'].parameters()).device
-    
-    if not isinstance(text, list): 
-        text = [text]
-    
-    assert len(text), 'text should not be empty'
-
-    text = [f"[Sbreak]{i}[Pbreak]{prompt}" for i in text]
-    text_token = models['tokenizer'](text, return_tensors='pt', add_special_tokens=False, padding=True).to(device)
-    text_mask = torch.ones(text_token['input_ids'].shape, dtype=bool, device=device)
-
-    inputs = {
-        'input_ids': text_token['input_ids'][...,None].expand(-1, -1, models['gpt'].num_vq),
-        'text_mask': text_mask,
-        'attention_mask': text_token['attention_mask'],
-    }
-    
-    LogitsWarpers = []
-    if top_P is not None:
-        LogitsWarpers.append(TopPLogitsWarper(top_P, min_tokens_to_keep=3))
-    if top_K is not None:
-        LogitsWarpers.append(TopKLogitsWarper(top_K, min_tokens_to_keep=3))
-        
-    LogitsProcessors = []
-    if repetition_penalty is not None and repetition_penalty != 1:
-        LogitsProcessors.append(CustomRepetitionPenaltyLogitsProcessorRepeat(repetition_penalty, len(models['tokenizer']), 16))
-    
-    result = models['gpt'].generate(
-        models['gpt'].get_emb(**inputs), inputs['input_ids'], 
-        temperature = torch.tensor([temperature,], device=device), 
-        attention_mask = inputs['attention_mask'],
-        LogitsWarpers = LogitsWarpers,
-        LogitsProcessors = LogitsProcessors,
-        eos_token = torch.tensor(models['tokenizer'].convert_tokens_to_ids('[Ebreak]'), device=device)[None], 
-        max_new_token = max_new_token, 
-        infer_text = True,
-        **kwargs
-    )
-    return result