import torch
## 迁移gpt模型的embedding层,插入两个新的token
gpt_dict = torch.load(r"GPT_SoVITS/pretrained_models/s1v3.ckpt",map_location=torch.device('cpu'))
gpt_shape = gpt_dict["weight"]["model.ar_text_embedding.word_embeddings.weight"].shape
print(gpt_shape)
# 一共512列，embedding_dim=512
first_part = gpt_dict["weight"]["model.ar_text_embedding.word_embeddings.weight"][:,:]
new_weight = torch.cat([first_part, torch.randn(21, 512)],dim=0)  # 由于embedding_dim=512，len(rus_symbols)=21，所以需要增加21列
gpt_dict["weight"]["model.ar_text_embedding.word_embeddings.weight"] = new_weight
gpt_dict["config"]["model"]["phoneme_vocab_size"]=753  # 原本是732，增加21个符号
torch.save(gpt_dict,r"GPT_SoVITS/pretrained_models/s1v3-ru.ckpt") # 保存新的模型


import torch
sovits_dict = torch.load(r"GPT_SoVITS/pretrained_models/v2Pro/s2Gv2ProPlus.pth",map_location=torch.device('cpu')) # 载入SoVITS模型
sovits_shape = sovits_dict["weight"]["enc_p.text_embedding.weight"].shape
print(sovits_shape)
first_part = sovits_dict["weight"]["enc_p.text_embedding.weight"][:,:]
new_weight = torch.cat([first_part,torch.randn(21,192)],dim=0) # 由于embedding_dim=192，len(rus_symbols)=21，所以需要增加21列
sovits_dict["weight"]["enc_p.text_embedding.weight"] = new_weight
torch.save(sovits_dict,r"GPT_SoVITS/pretrained_models/v2Pro/s2Gv2ProPlus-rus.pth") # 保存新的模型

print("success!")

# -*- coding: utf-8 -*- 

# 实现俄语的g2p
import epitran

epi = epitran.Epitran('rus-Cyrl',ligatures=True,tones=True)

def g2p(text):
    phones = epi.xsampa_list(text,normpunc=True)
    phones = ["R"+post_replace_rus(i) for i in phones if i !="?"]
    return phones

# 一些奇奇怪怪的符号替换成纯英文的符号,防止写到文件后解析出错
def post_replace_rus(rus):
    rep_map = {
        "：": ",",
        "；": ",",
        "，": ",",
        "。": ".",
        "！": "!",
        "？": "?",
        "\n": ".",
        "·": ",",
        "、": ",",
        "...": "…",
        "?":"",
    }

    if rus in rep_map.keys():
        rus = rep_map[rus]
    return rus.upper()

# 处理俄语的音素符号 去重清洗
def process_russian_symbols(file_path):
    import pandas as pd
        # 加载指定目录下的tsv文件
    df = pd.read_csv(file_path, sep='\t')  # 使用 tab 分隔符读取 tsv 文件
    
    # 读取 sentence 列
    sentences = df['sentence'].tolist()
    
    # 调用 g2p 函数处理每个句子
    # 这里假设 g2p 是一个已定义的函数，返回音素列表
    phonemes = []
    for sentence in sentences:
        phoneme_list = g2p(sentence)  # 假设 g2p 返回一个列表
        # if "RS\\:'" in phoneme_list:
        #     print(sentence)
        phonemes.extend(phoneme_list)
    
    # 去重并排序
    unique_sorted_phonemes = sorted(list(set(phonemes)))
    
    return unique_sorted_phonemes

rus_symbols={
    'RA',
    'RB',
    'RD', 
    'RE', 
    'RF', 
    'RG', 
    'RI', 
    'RJ', 
    'RK', 
    'RL', 
    'RM', 
    'RN', 
    'RO', 
    'RP', 
    'RR', 
    'RS', 
    'RT', 
    'RU', 
    'RV', 
    'RX', 
    'RZ'
}