mirror of
https://github.com/xming521/WeClone.git
synced 2026-08-28 18:07:28 +08:00
test(PII) add tests for PII
This commit is contained in:
@@ -25,6 +25,7 @@
|
||||
"1234567890",
|
||||
"hh"
|
||||
],
|
||||
"language": "en",
|
||||
"max_image_num": 2, // 单条数据最大图片数量
|
||||
"single_combine_strategy": "time_window", // 单人组成单句策略
|
||||
"qa_match_strategy": "time_window", // 组成qa策略
|
||||
|
||||
@@ -25,6 +25,7 @@
|
||||
"1234567890",
|
||||
"hh"
|
||||
],
|
||||
"language": "zh",
|
||||
"max_image_num": 2, // 单条数据最大图片数量
|
||||
"single_combine_strategy": "time_window", // 单人组成单句策略
|
||||
"qa_match_strategy": "time_window", // 组成qa策略
|
||||
|
||||
@@ -0,0 +1,68 @@
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from typing import cast
|
||||
|
||||
import pytest
|
||||
|
||||
# Import common functions from test_full_pipe
|
||||
from tests.test_full_pipe import (
|
||||
DATASET_CSV_DIR,
|
||||
PROJECT_ROOT_DIR,
|
||||
get_config_files,
|
||||
load_config_with_path,
|
||||
print_test_header,
|
||||
run_cli_command,
|
||||
setup_data_environment,
|
||||
test_logger,
|
||||
)
|
||||
from weclone.utils.config import load_config
|
||||
from weclone.utils.config_models import DataModality, WCMakeDatasetConfig
|
||||
from weclone.utils.log import logger
|
||||
|
||||
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
|
||||
|
||||
# Setup paths
|
||||
TESTS_DIR = os.path.dirname(__file__)
|
||||
TEST_DATA_PII_DIR = os.path.join(TESTS_DIR, "tests_data", "test_PII")
|
||||
|
||||
@pytest.mark.parametrize("config_file", get_config_files())
|
||||
def test_PII_make_dataset(config_file):
|
||||
"""Test PII data make-dataset functionality"""
|
||||
print_test_header("PII make-dataset", config_file)
|
||||
|
||||
setup_data_environment("test_PII")
|
||||
|
||||
# Load config and handle images if needed
|
||||
config: WCMakeDatasetConfig = cast(WCMakeDatasetConfig, load_config_with_path(config_file, "make_dataset"))
|
||||
|
||||
# Run make-dataset command
|
||||
result = run_cli_command(["make-dataset"], config_file)
|
||||
assert result.returncode == 0, f"make-dataset command execution failed for config {config_file}"
|
||||
|
||||
# Print all user messages from the dataset file with PII warning
|
||||
import json
|
||||
sft_file_path = os.path.join(PROJECT_ROOT_DIR, "dataset", "res_csv", "sft", "sft-my.json")
|
||||
if os.path.exists(sft_file_path):
|
||||
logger.warning("⚠️ WARNING: The following content contains unfiltered PII (Personally Identifiable Information):")
|
||||
logger.warning("=" * 80)
|
||||
|
||||
with open(sft_file_path, 'r', encoding='utf-8') as f:
|
||||
data = json.load(f)
|
||||
|
||||
for entry in data:
|
||||
if 'messages' in entry:
|
||||
for message in entry['messages']:
|
||||
if message.get('role') == 'user':
|
||||
logger.warning(f"User content: {message.get('content', '')}")
|
||||
|
||||
logger.warning("=" * 80)
|
||||
logger.warning("⚠️ END OF UNFILTERED PII CONTENT")
|
||||
|
||||
test_logger.info(f"✅ PII make-dataset test passed for config {config_file}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
# If running directly, run tests for all configs
|
||||
for config_file in get_config_files():
|
||||
test_PII_make_dataset(config_file)
|
||||
+22
-10
@@ -4,7 +4,7 @@ import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from typing import Optional, Union, cast
|
||||
from typing import Callable, Optional, Union, cast
|
||||
from unittest import mock
|
||||
|
||||
import pytest
|
||||
@@ -20,6 +20,7 @@ DATASET_CSV_DIR = os.path.join(PROJECT_ROOT, "dataset", "csv")
|
||||
TESTS_DIR = os.path.dirname(__file__)
|
||||
TEST_DATA_PERSON_DIR = os.path.join(TESTS_DIR, "tests_data", "test_person")
|
||||
|
||||
|
||||
# Backup directories
|
||||
BACKUP_DIR = os.path.join(PROJECT_ROOT, "test_backup")
|
||||
MODEL_OUTPUT_BACKUP = os.path.join(BACKUP_DIR, "model_output")
|
||||
@@ -67,10 +68,9 @@ def print_config_header(config_file: str):
|
||||
test_logger.info(" " * padding_left + title + " " * padding_right)
|
||||
test_logger.info("═" * line_length)
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def setup_test_environment():
|
||||
"""Setup test environment once for the entire test session"""
|
||||
test_logger.info("🔧 开始设置测试环境...")
|
||||
def setup_data_environment(data_folder_name: str = "test_person"):
|
||||
"""Setup test data environment for specified folder"""
|
||||
test_logger.info(f"🔧 设置 {data_folder_name} 测试数据...")
|
||||
|
||||
# Create backup directory
|
||||
if os.path.exists(BACKUP_DIR):
|
||||
@@ -89,15 +89,27 @@ def setup_test_environment():
|
||||
|
||||
os.makedirs(DATASET_CSV_DIR)
|
||||
|
||||
test_person_csv_dir = os.path.join(DATASET_CSV_DIR, "test_person")
|
||||
os.makedirs(test_person_csv_dir)
|
||||
# Setup specified test data folder
|
||||
test_data_source_dir = os.path.join(TESTS_DIR, "tests_data", data_folder_name)
|
||||
test_data_csv_dir = os.path.join(DATASET_CSV_DIR, data_folder_name)
|
||||
os.makedirs(test_data_csv_dir)
|
||||
|
||||
for item_name in os.listdir(TEST_DATA_PERSON_DIR):
|
||||
source_item_path = os.path.join(TEST_DATA_PERSON_DIR, item_name)
|
||||
for item_name in os.listdir(test_data_source_dir):
|
||||
source_item_path = os.path.join(test_data_source_dir, item_name)
|
||||
if os.path.isfile(source_item_path) and item_name.lower().endswith('.csv'):
|
||||
destination_item_path = os.path.join(test_person_csv_dir, item_name)
|
||||
destination_item_path = os.path.join(test_data_csv_dir, item_name)
|
||||
shutil.copy2(source_item_path, destination_item_path)
|
||||
|
||||
test_logger.info(f"✅ {data_folder_name} 测试数据设置完成")
|
||||
|
||||
@pytest.fixture(scope="session", autouse=True)
|
||||
def setup_test_environment():
|
||||
"""Setup test environment once for the entire test session"""
|
||||
test_logger.info("🔧 开始设置测试环境...")
|
||||
|
||||
# Use the generic setup function with default test_person data
|
||||
setup_data_environment("test_person")
|
||||
|
||||
test_logger.info("✅ 测试环境设置完成")
|
||||
|
||||
yield # This is where the testing happens
|
||||
|
||||
@@ -1,134 +0,0 @@
|
||||
# WeClone-audio 模块
|
||||
|
||||
WeClone-audio 是一个使用微信语音消息克隆声音的模块,使用模型实现高质量语音合成。
|
||||
### 显存需求
|
||||
**Spark-TTS** 推荐
|
||||
- **0.5B 模型**: 约 4GB 显存
|
||||
|
||||
**Llasa** (已弃用)
|
||||
- **3B 模型**: 约 16GB 显存
|
||||
- **1B 模型**: 约 9GB 显存
|
||||
|
||||
|
||||
|
||||
|
||||
## 1. 导出微信语音数据
|
||||
|
||||
### 1.1 准备工作
|
||||
- 使用 [PyWxDump](https://github.com/xaoyaoo/PyWxDump) 提取微信聊天记录
|
||||
- 下载软件并解密数据库
|
||||
- 点击聊天备份,导出类型选择"解密文件"
|
||||
|
||||
### 1.2 环境配置
|
||||
语音导出仅支持Windows环境
|
||||
WeClone Audio使用uv作为包管理器。
|
||||
```bash
|
||||
# 为 PyWxDump 创建 Python 环境和安装依赖
|
||||
#
|
||||
uv venv .venv-wx --python=3.10
|
||||
.venv-wx\Scripts\activate
|
||||
uv pip install pywxdump
|
||||
```
|
||||
|
||||
### 1.3 导出语音文件
|
||||
```bash
|
||||
python weclone-audio/src/get_sample_audio.py --db-path "导出数据库路径" --MsgSvrID "导出聊天记录的MsgSvrID字段"
|
||||
```
|
||||
|
||||
## 2. 语音合成推理
|
||||
### Spark-TTS模型
|
||||
|
||||
**环境安装**
|
||||
可不创建新环境,直接安装`sparktts`依赖组到WeClone共主环境
|
||||
|
||||
```bash
|
||||
uv venv .venv-sparktts --python=3.10
|
||||
source .venv-sparktts/bin/activate
|
||||
uv pip install --group sparktts -e .
|
||||
|
||||
git clone https://github.com/SparkAudio/Spark-TTS.git weclone-audio/src/Spark-TTS
|
||||
```
|
||||
|
||||
|
||||
**模型下载**
|
||||
|
||||
通过python下载:
|
||||
```python
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
# 假设此 Python 代码在 weclone-audio 目录下运行 模型将下载到 weclone-audio/pretrained_models/Spark-TTS-0.5B
|
||||
snapshot_download("SparkAudio/Spark-TTS-0.5B", local_dir="pretrained_models/Spark-TTS-0.5B")
|
||||
```
|
||||
|
||||
或通过git下载:
|
||||
```bash
|
||||
# 假设当前在 weclone-audio 目录
|
||||
mkdir -p pretrained_models
|
||||
|
||||
# Make sure you have git-lfs installed (https://git-lfs.com)
|
||||
git lfs install
|
||||
git clone https://huggingface.co/SparkAudio/Spark-TTS-0.5B pretrained_models/Spark-TTS-0.5B
|
||||
```
|
||||
使用代码推理
|
||||
```python
|
||||
import os
|
||||
import SparkTTS
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from SparkTTS import SparkTTS
|
||||
|
||||
# 假设此 Python 代码在 weclone-audio 目录下运行
|
||||
# 模型路径相对于当前目录
|
||||
model_path = "pretrained_models/Spark-TTS-0.5B"
|
||||
sample_audio = "sample.wav"
|
||||
output_audio = "output.wav"
|
||||
|
||||
model = SparkTTS(model_path, "cuda")
|
||||
|
||||
with torch.no_grad():
|
||||
wav = model.inference(
|
||||
text="晚上好啊,小可爱们,该睡觉了哦",
|
||||
prompt_speech_path=sample_audio, # 使用相对路径
|
||||
prompt_text="对,这就是我万人敬仰的太乙真人,虽然有点婴儿肥,但也掩不住我逼人的帅气。",
|
||||
)
|
||||
sf.write(output_audio, wav, samplerate=16000) # 使用相对路径
|
||||
```
|
||||
### Llasa模型 (已弃用)
|
||||
### 2.1 环境配置
|
||||
```bash
|
||||
# 创建并配置推理环境
|
||||
## 可不创建新环境,与LLaMA-Factory环境共用
|
||||
uv venv .venv-xcodec --python=3.9
|
||||
source .venv-xcodec/bin/activate
|
||||
uv pip install --group xcodec -e .
|
||||
# 退出环境
|
||||
deactivate
|
||||
|
||||
# 系统依赖安装(如果需要)
|
||||
sudo apt install python3-dev
|
||||
sudo apt install build-essential
|
||||
```
|
||||
|
||||
### 2.2 使用代码推理
|
||||
如果遇到问题,请尝试将参考音频转换为WAV或MP3格式,将其裁剪至15秒以内,并缩短提示文本。
|
||||
```python
|
||||
import os
|
||||
import soundfile as sf
|
||||
# 假设 text_to_speech.py 位于 src/ 或其他可导入的位置
|
||||
from text_to_speech import TextToSpeech
|
||||
|
||||
|
||||
sample_audio_text = "对,这就是我万人敬仰的太乙真人,虽然有点婴儿肥,但也掩不住我逼人的帅气。" # 示例音频文本
|
||||
# 假设此 Python 代码在 weclone-audio 目录下运行
|
||||
# 示例音频路径相对于当前目录
|
||||
sample_audio_path = "sample.wav"
|
||||
output_audio = "output.wav"
|
||||
|
||||
|
||||
tts = TextToSpeech(sample_audio_path, sample_audio_text)
|
||||
target_text = "晚上好啊" # 生成目标文本
|
||||
result = tts.infer(target_text)
|
||||
sf.write(output_audio, result[1], result[0]) # 使用相对路径
|
||||
```
|
||||
|
||||
@@ -1,11 +0,0 @@
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
from text_to_speech import TextToSpeech
|
||||
|
||||
sample_audio_text = "对,这就是我万人敬仰的太乙真人,虽然有点婴儿肥,但也掩不住我逼人的帅气。" # 示例音频文本
|
||||
sample_audio_path = os.path.join(os.path.dirname(__file__), "sample.wav") # 示例音频路径
|
||||
tts = TextToSpeech(sample_audio_path, sample_audio_text)
|
||||
target_text = "晚上好啊" # 生成目标文本
|
||||
result = tts.infer(target_text)
|
||||
sf.write(os.path.join(os.path.dirname(__file__), "output.wav"), result[1], result[0]) # 保存生成音频
|
||||
@@ -1,125 +0,0 @@
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
import torchaudio
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
from xcodec2.modeling_xcodec2 import XCodec2Model
|
||||
|
||||
|
||||
class TextToSpeech:
|
||||
def __init__(self, sample_audio_path, sample_audio_text):
|
||||
self.sample_audio_text = sample_audio_text
|
||||
# 初始化模型
|
||||
llasa_3b = "HKUSTAudio/Llasa-3B"
|
||||
xcodec2 = "HKUSTAudio/xcodec2"
|
||||
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(llasa_3b)
|
||||
self.llasa_3b_model = AutoModelForCausalLM.from_pretrained(
|
||||
llasa_3b,
|
||||
trust_remote_code=True,
|
||||
device_map="auto",
|
||||
)
|
||||
self.llasa_3b_model.eval()
|
||||
|
||||
self.xcodec_model = XCodec2Model.from_pretrained(xcodec2)
|
||||
self.xcodec_model.eval().cuda()
|
||||
|
||||
# 处理音频
|
||||
waveform, sample_rate = torchaudio.load(sample_audio_path)
|
||||
if len(waveform[0]) / sample_rate > 15:
|
||||
print("已将音频裁剪至前15秒。")
|
||||
waveform = waveform[:, : sample_rate * 15]
|
||||
|
||||
# 检查音频是否为立体声
|
||||
if waveform.size(0) > 1:
|
||||
waveform_mono = torch.mean(waveform, dim=0, keepdim=True)
|
||||
else:
|
||||
waveform_mono = waveform
|
||||
|
||||
self.prompt_wav = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)(waveform_mono)
|
||||
|
||||
# Encode the prompt wav
|
||||
vq_code_prompt = self.xcodec_model.encode_code(input_waveform=self.prompt_wav)
|
||||
vq_code_prompt = vq_code_prompt[0, 0, :]
|
||||
self.speech_ids_prefix = self.ids_to_speech_tokens(vq_code_prompt)
|
||||
self.speech_end_id = self.tokenizer.convert_tokens_to_ids("<|SPEECH_GENERATION_END|>")
|
||||
|
||||
def ids_to_speech_tokens(self, speech_ids):
|
||||
speech_tokens_str = []
|
||||
for speech_id in speech_ids:
|
||||
speech_tokens_str.append(f"<|s_{speech_id}|>")
|
||||
return speech_tokens_str
|
||||
|
||||
def extract_speech_ids(self, speech_tokens_str):
|
||||
speech_ids = []
|
||||
for token_str in speech_tokens_str:
|
||||
if token_str.startswith("<|s_") and token_str.endswith("|>"):
|
||||
num_str = token_str[4:-2]
|
||||
num = int(num_str)
|
||||
speech_ids.append(num)
|
||||
else:
|
||||
print(f"Unexpected token: {token_str}")
|
||||
return speech_ids
|
||||
|
||||
@torch.inference_mode()
|
||||
def infer(self, target_text):
|
||||
if len(target_text) == 0:
|
||||
return None
|
||||
elif len(target_text) > 300:
|
||||
print("文本过长,请保持在300字符以内。")
|
||||
target_text = target_text[:300]
|
||||
|
||||
input_text = self.sample_audio_text + " " + target_text
|
||||
|
||||
formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
|
||||
|
||||
chat = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Convert the text to speech:" + formatted_text,
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "<|SPEECH_GENERATION_START|>" + "".join(self.speech_ids_prefix),
|
||||
},
|
||||
]
|
||||
|
||||
input_ids = self.tokenizer.apply_chat_template(
|
||||
chat, tokenize=True, return_tensors="pt", continue_final_message=True
|
||||
)
|
||||
input_ids = input_ids.to("cuda")
|
||||
|
||||
outputs = self.llasa_3b_model.generate(
|
||||
input_ids,
|
||||
max_length=2048,
|
||||
eos_token_id=self.speech_end_id,
|
||||
do_sample=True,
|
||||
top_p=1,
|
||||
temperature=0.8,
|
||||
)
|
||||
generated_ids = outputs[0][input_ids.shape[1] - len(self.speech_ids_prefix) : -1]
|
||||
|
||||
speech_tokens = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
|
||||
|
||||
speech_tokens = self.extract_speech_ids(speech_tokens)
|
||||
speech_tokens = torch.tensor(speech_tokens).cuda().unsqueeze(0).unsqueeze(0)
|
||||
|
||||
gen_wav = self.xcodec_model.decode_code(speech_tokens)
|
||||
gen_wav = gen_wav[:, :, self.prompt_wav.shape[1] :]
|
||||
|
||||
return (16000, gen_wav[0, 0, :].cpu().numpy())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
# 如果遇到问题,请尝试将参考音频转换为WAV或MP3格式,将其裁剪至15秒以内,并缩短提示文本。
|
||||
sample_audio_text = "对,这就是我万人敬仰的太乙真人,虽然有点婴儿肥,但也掩不住我逼人的帅气。"
|
||||
sample_audio_path = os.path.join(os.path.dirname(__file__), "sample.wav")
|
||||
|
||||
tts = TextToSpeech(sample_audio_path, sample_audio_text)
|
||||
target_text = "晚上好啊,吃了吗您"
|
||||
result = tts.infer(target_text)
|
||||
sf.write(os.path.join(os.path.dirname(__file__), "output.wav"), result[1], result[0])
|
||||
target_text = "我是老北京正黄旗!"
|
||||
result = tts.infer(target_text)
|
||||
sf.write(os.path.join(os.path.dirname(__file__), "output1.wav"), result[1], result[0])
|
||||
Submodule weclone-audio/src/Spark-TTS deleted from ee29f36806
@@ -1,215 +0,0 @@
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Tuple
|
||||
|
||||
import torch
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||||
|
||||
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "./Spark-TTS")))
|
||||
from sparktts.models.audio_tokenizer import BiCodecTokenizer
|
||||
from sparktts.utils.file import load_config
|
||||
from sparktts.utils.token_parser import GENDER_MAP, LEVELS_MAP, TASK_TOKEN_MAP
|
||||
|
||||
|
||||
class SparkTTS:
|
||||
"""
|
||||
Spark-TTS for text-to-speech generation.
|
||||
"""
|
||||
|
||||
def __init__(self, model_dir: Path, device: torch.device = torch.device("cuda:0")):
|
||||
"""
|
||||
Initializes the SparkTTS model with the provided configurations and device.
|
||||
|
||||
Args:
|
||||
model_dir (Path): Directory containing the model and config files.
|
||||
device (torch.device): The device (CPU/GPU) to run the model on.
|
||||
"""
|
||||
self.device = device
|
||||
self.model_dir = model_dir
|
||||
self.configs = load_config(f"{model_dir}/config.yaml")
|
||||
self.sample_rate = self.configs["sample_rate"]
|
||||
self._initialize_inference()
|
||||
|
||||
def _initialize_inference(self):
|
||||
"""Initializes the tokenizer, model, and audio tokenizer for inference."""
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(f"{self.model_dir}/LLM")
|
||||
self.model = AutoModelForCausalLM.from_pretrained(f"{self.model_dir}/LLM")
|
||||
self.audio_tokenizer = BiCodecTokenizer(self.model_dir, device=self.device)
|
||||
self.model.to(self.device)
|
||||
|
||||
def process_prompt(
|
||||
self,
|
||||
text: str,
|
||||
prompt_speech_path: Path,
|
||||
prompt_text: str = None,
|
||||
) -> Tuple[str, torch.Tensor]:
|
||||
"""
|
||||
Process input for voice cloning.
|
||||
|
||||
Args:
|
||||
text (str): The text input to be converted to speech.
|
||||
prompt_speech_path (Path): Path to the audio file used as a prompt.
|
||||
prompt_text (str, optional): Transcript of the prompt audio.
|
||||
|
||||
Return:
|
||||
Tuple[str, torch.Tensor]: Input prompt; global tokens
|
||||
"""
|
||||
|
||||
global_token_ids, semantic_token_ids = self.audio_tokenizer.tokenize(prompt_speech_path)
|
||||
global_tokens = "".join([f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()])
|
||||
|
||||
# Prepare the input tokens for the model
|
||||
if prompt_text is not None:
|
||||
semantic_tokens = "".join([f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()])
|
||||
inputs = [
|
||||
TASK_TOKEN_MAP["tts"],
|
||||
"<|start_content|>",
|
||||
prompt_text,
|
||||
text,
|
||||
"<|end_content|>",
|
||||
"<|start_global_token|>",
|
||||
global_tokens,
|
||||
"<|end_global_token|>",
|
||||
"<|start_semantic_token|>",
|
||||
semantic_tokens,
|
||||
]
|
||||
else:
|
||||
inputs = [
|
||||
TASK_TOKEN_MAP["tts"],
|
||||
"<|start_content|>",
|
||||
text,
|
||||
"<|end_content|>",
|
||||
"<|start_global_token|>",
|
||||
global_tokens,
|
||||
"<|end_global_token|>",
|
||||
]
|
||||
|
||||
inputs = "".join(inputs)
|
||||
|
||||
return inputs, global_token_ids
|
||||
|
||||
def process_prompt_control(
|
||||
self,
|
||||
gender: str,
|
||||
pitch: str,
|
||||
speed: str,
|
||||
text: str,
|
||||
):
|
||||
"""
|
||||
Process input for voice creation.
|
||||
|
||||
Args:
|
||||
gender (str): female | male.
|
||||
pitch (str): very_low | low | moderate | high | very_high
|
||||
speed (str): very_low | low | moderate | high | very_high
|
||||
text (str): The text input to be converted to speech.
|
||||
|
||||
Return:
|
||||
str: Input prompt
|
||||
"""
|
||||
assert gender in GENDER_MAP.keys()
|
||||
assert pitch in LEVELS_MAP.keys()
|
||||
assert speed in LEVELS_MAP.keys()
|
||||
|
||||
gender_id = GENDER_MAP[gender]
|
||||
pitch_level_id = LEVELS_MAP[pitch]
|
||||
speed_level_id = LEVELS_MAP[speed]
|
||||
|
||||
pitch_label_tokens = f"<|pitch_label_{pitch_level_id}|>"
|
||||
speed_label_tokens = f"<|speed_label_{speed_level_id}|>"
|
||||
gender_tokens = f"<|gender_{gender_id}|>"
|
||||
|
||||
attribte_tokens = "".join([gender_tokens, pitch_label_tokens, speed_label_tokens])
|
||||
|
||||
control_tts_inputs = [
|
||||
TASK_TOKEN_MAP["controllable_tts"],
|
||||
"<|start_content|>",
|
||||
text,
|
||||
"<|end_content|>",
|
||||
"<|start_style_label|>",
|
||||
attribte_tokens,
|
||||
"<|end_style_label|>",
|
||||
]
|
||||
|
||||
return "".join(control_tts_inputs)
|
||||
|
||||
@torch.no_grad()
|
||||
def inference(
|
||||
self,
|
||||
text: str,
|
||||
prompt_speech_path: Path = None,
|
||||
prompt_text: str = None,
|
||||
gender: str = None,
|
||||
pitch: str = None,
|
||||
speed: str = None,
|
||||
temperature: float = 0.8,
|
||||
top_k: float = 50,
|
||||
top_p: float = 0.95,
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Performs inference to generate speech from text, incorporating prompt audio and/or text.
|
||||
|
||||
Args:
|
||||
text (str): The text input to be converted to speech.
|
||||
prompt_speech_path (Path): Path to the audio file used as a prompt.
|
||||
prompt_text (str, optional): Transcript of the prompt audio.
|
||||
gender (str): female | male.
|
||||
pitch (str): very_low | low | moderate | high | very_high
|
||||
speed (str): very_low | low | moderate | high | very_high
|
||||
temperature (float, optional): Sampling temperature for controlling randomness. Default is 0.8.
|
||||
top_k (float, optional): Top-k sampling parameter. Default is 50.
|
||||
top_p (float, optional): Top-p (nucleus) sampling parameter. Default is 0.95.
|
||||
|
||||
Returns:
|
||||
torch.Tensor: Generated waveform as a tensor.
|
||||
"""
|
||||
if gender is not None:
|
||||
prompt = self.process_prompt_control(gender, pitch, speed, text)
|
||||
|
||||
else:
|
||||
prompt, global_token_ids = self.process_prompt(text, prompt_speech_path, prompt_text)
|
||||
model_inputs = self.tokenizer([prompt], return_tensors="pt").to(self.device)
|
||||
|
||||
# Generate speech using the model
|
||||
generated_ids = self.model.generate(
|
||||
**model_inputs,
|
||||
max_new_tokens=3000,
|
||||
do_sample=True,
|
||||
top_k=top_k,
|
||||
top_p=top_p,
|
||||
temperature=temperature,
|
||||
)
|
||||
|
||||
# Trim the output tokens to remove the input tokens
|
||||
generated_ids = [
|
||||
output_ids[len(input_ids) :]
|
||||
for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
|
||||
]
|
||||
|
||||
# Decode the generated tokens into text
|
||||
predicts = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
|
||||
|
||||
# Extract semantic token IDs from the generated text
|
||||
pred_semantic_ids = (
|
||||
torch.tensor([int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicts)])
|
||||
.long()
|
||||
.unsqueeze(0)
|
||||
)
|
||||
|
||||
if gender is not None:
|
||||
global_token_ids = (
|
||||
torch.tensor([int(token) for token in re.findall(r"bicodec_global_(\d+)", predicts)])
|
||||
.long()
|
||||
.unsqueeze(0)
|
||||
.unsqueeze(0)
|
||||
)
|
||||
|
||||
# Convert semantic tokens back to waveform
|
||||
wav = self.audio_tokenizer.detokenize(
|
||||
global_token_ids.to(self.device).squeeze(0),
|
||||
pred_semantic_ids.to(self.device),
|
||||
)
|
||||
|
||||
return wav
|
||||
@@ -1,40 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from pywxdump.db import MediaHandler
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Extract audio from WeChat database")
|
||||
parser.add_argument("--db-path", type=str, required=True, help="Path to WeChat database file")
|
||||
parser.add_argument("--MsgSvrID", type=str, required=True, help="Message server ID of the audio")
|
||||
parser.add_argument(
|
||||
"--save-path",
|
||||
type=str,
|
||||
default=os.path.join(os.path.dirname(__file__), "sample.wav"),
|
||||
help="Path to save the audio file (default: sample.wav in script directory)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--rate", type=int, default=24000, help="Sample rate for audio conversion (default: 24000)"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
config = {
|
||||
"key": "test1",
|
||||
"type": "sqlite",
|
||||
"path": args.db_path,
|
||||
}
|
||||
|
||||
t1 = MediaHandler(config)
|
||||
t1.get_audio(
|
||||
MsgSvrID=args.MsgSvrID,
|
||||
is_play=True,
|
||||
is_wave=True,
|
||||
save_path=args.save_path,
|
||||
rate=args.rate,
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,17 +0,0 @@
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from SparkTTS import SparkTTS
|
||||
|
||||
model = SparkTTS("weclone-audio/pretrained_models/Spark-TTS-0.5B", "cuda")
|
||||
|
||||
|
||||
with torch.no_grad():
|
||||
wav = model.inference(
|
||||
text="晚上好啊,小可爱们,该睡觉了哦",
|
||||
prompt_speech_path=os.path.join(os.path.dirname(__file__), "sample.wav"),
|
||||
prompt_text="对,这就是我万人敬仰的太乙真人,虽然有点婴儿肥,但也掩不住我逼人的帅气。",
|
||||
)
|
||||
sf.write(os.path.join(os.path.dirname(__file__), "output.wav"), wav, samplerate=16000)
|
||||
print("生成成功!")
|
||||
Binary file not shown.
@@ -1,14 +0,0 @@
|
||||
API_KEY=your_api_key_here
|
||||
PORT=5050
|
||||
|
||||
DEFAULT_VOICE=en-US-AvaNeural
|
||||
DEFAULT_RESPONSE_FORMAT=mp3
|
||||
DEFAULT_SPEED=1.0
|
||||
|
||||
DEFAULT_LANGUAGE=en-US
|
||||
|
||||
REQUIRE_API_KEY=True
|
||||
|
||||
REMOVE_FILTER=False
|
||||
|
||||
EXPAND_API=True
|
||||
@@ -1,64 +0,0 @@
|
||||
import re
|
||||
|
||||
import emoji
|
||||
|
||||
|
||||
def prepare_tts_input_with_context(text: str) -> str:
|
||||
"""
|
||||
Prepares text for a TTS API by cleaning Markdown and adding minimal contextual hints
|
||||
for certain Markdown elements like headers. Preserves paragraph separation.
|
||||
|
||||
Args:
|
||||
text (str): The raw text containing Markdown or other formatting.
|
||||
|
||||
Returns:
|
||||
str: Cleaned text with contextual hints suitable for TTS input.
|
||||
"""
|
||||
|
||||
# Remove emojis
|
||||
text = emoji.replace_emoji(text, replace='')
|
||||
|
||||
# Add context for headers
|
||||
def header_replacer(match):
|
||||
level = len(match.group(1)) # Number of '#' symbols
|
||||
header_text = match.group(2).strip()
|
||||
if level == 1:
|
||||
return f"Title — {header_text}\n"
|
||||
elif level == 2:
|
||||
return f"Section — {header_text}\n"
|
||||
else:
|
||||
return f"Subsection — {header_text}\n"
|
||||
|
||||
text = re.sub(r"^(#{1,6})\s+(.*)", header_replacer, text, flags=re.MULTILINE)
|
||||
|
||||
# Announce links (currently commented out for potential future use)
|
||||
# text = re.sub(r"\[([^\]]+)\]\((https?:\/\/[^\)]+)\)", r"\1 (link: \2)", text)
|
||||
|
||||
# Remove links while keeping the link text
|
||||
text = re.sub(r"\[([^\]]+)\]\([^\)]+\)", r"\1", text)
|
||||
|
||||
# Describe inline code
|
||||
text = re.sub(r"`([^`]+)`", r"code snippet: \1", text)
|
||||
|
||||
# Remove bold/italic symbols but keep the content
|
||||
text = re.sub(r"(\*\*|__|\*|_)", '', text)
|
||||
|
||||
# Remove code blocks (multi-line) with a description
|
||||
text = re.sub(r"```([\s\S]+?)```", r"(code block omitted)", text)
|
||||
|
||||
# Remove image syntax but add alt text if available
|
||||
text = re.sub(r"!\[([^\]]*)\]\([^\)]+\)", r"Image: \1", text)
|
||||
|
||||
# Remove HTML tags
|
||||
text = re.sub(r"</?[^>]+(>|$)", '', text)
|
||||
|
||||
# Normalize line breaks
|
||||
text = re.sub(r"\n{2,}", '\n\n', text) # Ensure consistent paragraph separation
|
||||
|
||||
# Replace multiple spaces within lines
|
||||
text = re.sub(r" {2,}", ' ', text)
|
||||
|
||||
# Trim leading and trailing whitespace from the whole text
|
||||
text = text.strip()
|
||||
|
||||
return text
|
||||
@@ -1,5 +0,0 @@
|
||||
flask
|
||||
gevent
|
||||
python-dotenv
|
||||
edge-tts
|
||||
emoji
|
||||
@@ -1,167 +0,0 @@
|
||||
# server.py
|
||||
|
||||
import os
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from flask import Flask, jsonify, request, send_file
|
||||
from gevent.pywsgi import WSGIServer
|
||||
from handle_text import prepare_tts_input_with_context
|
||||
from tts_handler import generate_speech, get_models, get_voices
|
||||
from utils import AUDIO_FORMAT_MIME_TYPES, getenv_bool, require_api_key
|
||||
|
||||
app = Flask(__name__)
|
||||
load_dotenv()
|
||||
|
||||
API_KEY = os.getenv('API_KEY', 'your_api_key_here')
|
||||
PORT = int(os.getenv('PORT', 5050))
|
||||
|
||||
DEFAULT_VOICE = os.getenv('DEFAULT_VOICE', 'en-US-AvaNeural')
|
||||
DEFAULT_RESPONSE_FORMAT = os.getenv('DEFAULT_RESPONSE_FORMAT', 'mp3')
|
||||
DEFAULT_SPEED = float(os.getenv('DEFAULT_SPEED', 1.0))
|
||||
|
||||
REMOVE_FILTER = getenv_bool('REMOVE_FILTER', False)
|
||||
EXPAND_API = getenv_bool('EXPAND_API', True)
|
||||
|
||||
# DEFAULT_MODEL = os.getenv('DEFAULT_MODEL', 'tts-1')
|
||||
|
||||
@app.route('/v1/audio/speech', methods=['POST'])
|
||||
@app.route('/audio/speech', methods=['POST']) # Add this line for the alias
|
||||
@require_api_key
|
||||
def text_to_speech():
|
||||
data = request.json
|
||||
if not data or 'input' not in data:
|
||||
return jsonify({"error": "Missing 'input' in request body"}), 400
|
||||
|
||||
text = data.get('input')
|
||||
|
||||
if not REMOVE_FILTER:
|
||||
text = prepare_tts_input_with_context(text)
|
||||
|
||||
# model = data.get('model', DEFAULT_MODEL)
|
||||
voice = data.get('voice', DEFAULT_VOICE)
|
||||
|
||||
response_format = data.get('response_format', DEFAULT_RESPONSE_FORMAT)
|
||||
speed = float(data.get('speed', DEFAULT_SPEED))
|
||||
|
||||
mime_type = AUDIO_FORMAT_MIME_TYPES.get(response_format, "audio/mpeg")
|
||||
|
||||
# Generate the audio file in the specified format with speed adjustment
|
||||
output_file_path = generate_speech(text, voice, response_format, speed)
|
||||
|
||||
# Return the file with the correct MIME type
|
||||
return send_file(output_file_path, mimetype=mime_type, as_attachment=True, download_name=f"speech.{response_format}")
|
||||
|
||||
@app.route('/v1/models', methods=['GET', 'POST'])
|
||||
@app.route('/models', methods=['GET', 'POST'])
|
||||
@require_api_key
|
||||
def list_models():
|
||||
return jsonify({"data": get_models()})
|
||||
|
||||
@app.route('/v1/voices', methods=['GET', 'POST'])
|
||||
@app.route('/voices', methods=['GET', 'POST'])
|
||||
@require_api_key
|
||||
def list_voices():
|
||||
specific_language = None
|
||||
|
||||
data = request.args if request.method == 'GET' else request.json
|
||||
if data and ('language' in data or 'locale' in data):
|
||||
specific_language = data.get('language') if 'language' in data else data.get('locale')
|
||||
|
||||
return jsonify({"voices": get_voices(specific_language)})
|
||||
|
||||
@app.route('/v1/voices/all', methods=['GET', 'POST'])
|
||||
@app.route('/voices/all', methods=['GET', 'POST'])
|
||||
@require_api_key
|
||||
def list_all_voices():
|
||||
return jsonify({"voices": get_voices('all')})
|
||||
|
||||
"""
|
||||
Support for ElevenLabs and Azure AI Speech
|
||||
(currently in beta)
|
||||
"""
|
||||
|
||||
# http://localhost:5050/elevenlabs/v1/text-to-speech
|
||||
# http://localhost:5050/elevenlabs/v1/text-to-speech/en-US-AndrewNeural
|
||||
@app.route('/elevenlabs/v1/text-to-speech/<voice_id>', methods=['POST'])
|
||||
@require_api_key
|
||||
def elevenlabs_tts(voice_id):
|
||||
if not EXPAND_API:
|
||||
return jsonify({"error": f"Endpoint not allowed"}), 500
|
||||
|
||||
# Parse the incoming JSON payload
|
||||
try:
|
||||
payload = request.json
|
||||
if not payload or 'text' not in payload:
|
||||
return jsonify({"error": "Missing 'text' in request body"}), 400
|
||||
except Exception as e:
|
||||
return jsonify({"error": f"Invalid JSON payload: {str(e)}"}), 400
|
||||
|
||||
text = payload['text']
|
||||
|
||||
if not REMOVE_FILTER:
|
||||
text = prepare_tts_input_with_context(text)
|
||||
|
||||
voice = voice_id # ElevenLabs uses the voice_id in the URL
|
||||
|
||||
# Use default settings for edge-tts
|
||||
response_format = 'mp3'
|
||||
speed = DEFAULT_SPEED # Optional customization via payload.get('speed', DEFAULT_SPEED)
|
||||
|
||||
# Generate speech using edge-tts
|
||||
try:
|
||||
output_file_path = generate_speech(text, voice, response_format, speed)
|
||||
except Exception as e:
|
||||
return jsonify({"error": f"TTS generation failed: {str(e)}"}), 500
|
||||
|
||||
# Return the generated audio file
|
||||
return send_file(output_file_path, mimetype="audio/mpeg", as_attachment=True, download_name="speech.mp3")
|
||||
|
||||
# tts.speech.microsoft.com/cognitiveservices/v1
|
||||
# https://{region}.tts.speech.microsoft.com/cognitiveservices/v1
|
||||
# http://localhost:5050/azure/cognitiveservices/v1
|
||||
@app.route('/azure/cognitiveservices/v1', methods=['POST'])
|
||||
@require_api_key
|
||||
def azure_tts():
|
||||
if not EXPAND_API:
|
||||
return jsonify({"error": f"Endpoint not allowed"}), 500
|
||||
|
||||
# Parse the SSML payload
|
||||
try:
|
||||
ssml_data = request.data.decode('utf-8')
|
||||
if not ssml_data:
|
||||
return jsonify({"error": "Missing SSML payload"}), 400
|
||||
|
||||
# Extract the text and voice from SSML
|
||||
from xml.etree import ElementTree as ET
|
||||
root = ET.fromstring(ssml_data)
|
||||
text = root.find('.//{http://www.w3.org/2001/10/synthesis}voice').text
|
||||
voice = root.find('.//{http://www.w3.org/2001/10/synthesis}voice').get('name')
|
||||
except Exception as e:
|
||||
return jsonify({"error": f"Invalid SSML payload: {str(e)}"}), 400
|
||||
|
||||
# Use default settings for edge-tts
|
||||
response_format = 'mp3'
|
||||
speed = DEFAULT_SPEED
|
||||
|
||||
if not REMOVE_FILTER:
|
||||
text = prepare_tts_input_with_context(text)
|
||||
|
||||
# Generate speech using edge-tts
|
||||
try:
|
||||
output_file_path = generate_speech(text, voice, response_format, speed)
|
||||
except Exception as e:
|
||||
return jsonify({"error": f"TTS generation failed: {str(e)}"}), 500
|
||||
|
||||
# Return the generated audio file
|
||||
return send_file(output_file_path, mimetype="audio/mpeg", as_attachment=True, download_name="speech.mp3")
|
||||
|
||||
print(f" Edge TTS (Free Azure TTS) Replacement for OpenAI's TTS API")
|
||||
print(f" ")
|
||||
print(f" * Serving OpenAI Edge TTS")
|
||||
print(f" * Server running on http://localhost:{PORT}")
|
||||
print(f" * TTS Endpoint: http://localhost:{PORT}/v1/audio/speech")
|
||||
print(f" ")
|
||||
|
||||
if __name__ == '__main__':
|
||||
http_server = WSGIServer(('0.0.0.0', PORT), app)
|
||||
http_server.serve_forever()
|
||||
@@ -1,134 +0,0 @@
|
||||
import asyncio
|
||||
import os
|
||||
import subprocess
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
|
||||
import edge_tts
|
||||
|
||||
# Language default (environment variable)
|
||||
DEFAULT_LANGUAGE = os.getenv('DEFAULT_LANGUAGE', 'en-US')
|
||||
|
||||
# OpenAI voice names mapped to edge-tts equivalents
|
||||
voice_mapping = {
|
||||
'alloy': 'en-US-AvaNeural',
|
||||
'echo': 'en-US-AndrewNeural',
|
||||
'fable': 'en-GB-SoniaNeural',
|
||||
'onyx': 'en-US-EricNeural',
|
||||
'nova': 'en-US-SteffanNeural',
|
||||
'shimmer': 'en-US-EmmaNeural'
|
||||
}
|
||||
|
||||
def is_ffmpeg_installed():
|
||||
"""Check if FFmpeg is installed and accessible."""
|
||||
try:
|
||||
subprocess.run(['ffmpeg', '-version'], check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
return True
|
||||
except (subprocess.CalledProcessError, FileNotFoundError):
|
||||
return False
|
||||
|
||||
async def _generate_audio(text, voice, response_format, speed):
|
||||
"""Generate TTS audio and optionally convert to a different format."""
|
||||
# Determine if the voice is an OpenAI-compatible voice or a direct edge-tts voice
|
||||
edge_tts_voice = voice_mapping.get(voice, voice) # Use mapping if in OpenAI names, otherwise use as-is
|
||||
|
||||
# Generate the TTS output in mp3 format first
|
||||
temp_output_file = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
|
||||
|
||||
# Convert speed to SSML rate format
|
||||
try:
|
||||
speed_rate = speed_to_rate(speed) # Convert speed value to "+X%" or "-X%"
|
||||
except Exception as e:
|
||||
print(f"Error converting speed: {e}. Defaulting to +0%.")
|
||||
speed_rate = "+0%"
|
||||
|
||||
# Generate the MP3 file
|
||||
communicator = edge_tts.Communicate(text=text, voice=edge_tts_voice, rate=speed_rate)
|
||||
await communicator.save(temp_output_file.name)
|
||||
|
||||
# If the requested format is mp3, return the generated file directly
|
||||
if response_format == "mp3":
|
||||
return temp_output_file.name
|
||||
|
||||
# Check if FFmpeg is installed
|
||||
if not is_ffmpeg_installed():
|
||||
print("FFmpeg is not available. Returning unmodified mp3 file.")
|
||||
return temp_output_file.name
|
||||
|
||||
# Create a new temporary file for the converted output
|
||||
converted_output_file = tempfile.NamedTemporaryFile(delete=False, suffix=f".{response_format}")
|
||||
|
||||
# Build the FFmpeg command
|
||||
ffmpeg_command = [
|
||||
"ffmpeg",
|
||||
"-i", temp_output_file.name, # Input file
|
||||
"-c:a", {
|
||||
"aac": "aac",
|
||||
"mp3": "libmp3lame",
|
||||
"wav": "pcm_s16le",
|
||||
"opus": "libopus",
|
||||
"flac": "flac"
|
||||
}.get(response_format, "aac"), # Default to AAC if unknown
|
||||
"-b:a", "192k" if response_format != "wav" else None, # Bitrate not needed for WAV
|
||||
"-f", {
|
||||
"aac": "mp4", # AAC in MP4 container
|
||||
"mp3": "mp3",
|
||||
"wav": "wav",
|
||||
"opus": "ogg",
|
||||
"flac": "flac"
|
||||
}.get(response_format, response_format), # Default to matching format
|
||||
"-y", # Overwrite without prompt
|
||||
converted_output_file.name # Output file
|
||||
]
|
||||
|
||||
try:
|
||||
# Run FFmpeg command and ensure no errors occur
|
||||
subprocess.run(ffmpeg_command, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
|
||||
except subprocess.CalledProcessError as e:
|
||||
raise RuntimeError(f"FFmpeg error during audio conversion: {e}")
|
||||
|
||||
# Clean up the original temporary file
|
||||
Path(temp_output_file.name).unlink(missing_ok=True)
|
||||
|
||||
return converted_output_file.name
|
||||
|
||||
def generate_speech(text, voice, response_format, speed=1.0):
|
||||
return asyncio.run(_generate_audio(text, voice, response_format, speed))
|
||||
|
||||
def get_models():
|
||||
return [
|
||||
{"id": "tts-1", "name": "Text-to-speech v1"},
|
||||
{"id": "tts-1-hd", "name": "Text-to-speech v1 HD"}
|
||||
]
|
||||
|
||||
async def _get_voices(language=None):
|
||||
# List all voices, filter by language if specified
|
||||
all_voices = await edge_tts.list_voices()
|
||||
language = language or DEFAULT_LANGUAGE # Use default if no language specified
|
||||
filtered_voices = [
|
||||
{"name": v['ShortName'], "gender": v['Gender'], "language": v['Locale']}
|
||||
for v in all_voices if language == 'all' or language is None or v['Locale'] == language
|
||||
]
|
||||
return filtered_voices
|
||||
|
||||
def get_voices(language=None):
|
||||
return asyncio.run(_get_voices(language))
|
||||
|
||||
def speed_to_rate(speed: float) -> str:
|
||||
"""
|
||||
Converts a multiplicative speed value to the edge-tts "rate" format.
|
||||
|
||||
Args:
|
||||
speed (float): The multiplicative speed value (e.g., 1.5 for +50%, 0.5 for -50%).
|
||||
|
||||
Returns:
|
||||
str: The formatted "rate" string (e.g., "+50%" or "-50%").
|
||||
"""
|
||||
if speed < 0 or speed > 2:
|
||||
raise ValueError("Speed must be between 0 and 2 (inclusive).")
|
||||
|
||||
# Convert speed to percentage change
|
||||
percentage_change = (speed - 1) * 100
|
||||
|
||||
# Format with a leading "+" or "-" as required
|
||||
return f"{percentage_change:+.0f}%"
|
||||
@@ -1,39 +0,0 @@
|
||||
# utils.py
|
||||
|
||||
import os
|
||||
from functools import wraps
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from flask import jsonify, request
|
||||
|
||||
load_dotenv()
|
||||
|
||||
def getenv_bool(name: str, default: bool = False) -> bool:
|
||||
return os.getenv(name, str(default)).lower() in ("yes", "y", "true", "1", "t")
|
||||
|
||||
API_KEY = os.getenv('API_KEY', 'your_api_key_here')
|
||||
REQUIRE_API_KEY = getenv_bool('REQUIRE_API_KEY', True)
|
||||
|
||||
def require_api_key(f):
|
||||
@wraps(f)
|
||||
def decorated_function(*args, **kwargs):
|
||||
if not REQUIRE_API_KEY:
|
||||
return f(*args, **kwargs)
|
||||
auth_header = request.headers.get('Authorization')
|
||||
if not auth_header or not auth_header.startswith('Bearer '):
|
||||
return jsonify({"error": "Missing or invalid API key"}), 401
|
||||
token = auth_header.split('Bearer ')[1]
|
||||
if token != API_KEY:
|
||||
return jsonify({"error": "Invalid API key"}), 401
|
||||
return f(*args, **kwargs)
|
||||
return decorated_function
|
||||
|
||||
# Mapping of audio format to MIME type
|
||||
AUDIO_FORMAT_MIME_TYPES = {
|
||||
"mp3": "audio/mpeg",
|
||||
"opus": "audio/ogg",
|
||||
"aac": "audio/aac",
|
||||
"flac": "audio/flac",
|
||||
"wav": "audio/wav",
|
||||
"pcm": "audio/L16"
|
||||
}
|
||||
@@ -72,7 +72,7 @@ class PIIDetector:
|
||||
def _add_custom_recognizers(self, language: str):
|
||||
# Create numeric ID recognizer - matches 5+ digit numbers or numbers with - separators
|
||||
numeric_id_patterns = [
|
||||
Pattern(name="numeric_id", regex=r"\b(?:\d{5,}|\d+-\d+(?:-\d+)*)\b", score=0.8),
|
||||
Pattern(name="numeric_id", regex=r"\b(?:[A-Za-z]*\d{5,}[A-Za-z]*|\d+-\d+(?:-\d+)*)\b", score=0.8),
|
||||
Pattern(name="unicode_escape_id", regex=r"\\u[0-9a-fA-F]{4}", score=0.8),
|
||||
Pattern(name="hex_escape_id", regex=r"\\xa0", score=0.8),
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user