1308 lines
55 KiB
Python
1308 lines
55 KiB
Python
import random
|
||
import torch
|
||
import numpy as np
|
||
from typing import List, Dict, Optional, Tuple, Union
|
||
from tqdm import tqdm
|
||
import logging
|
||
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
||
import json
|
||
import requests
|
||
import time
|
||
import re
|
||
import os
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
def get_config():
|
||
"""Always use centralized config loader unless running standalone (no utils.config_loader)"""
|
||
try:
|
||
from utils.config_loader import get_config as _get_config
|
||
return _get_config()
|
||
except ImportError:
|
||
try:
|
||
from bge_finetune.utils.config_loader import get_config as _get_config # type: ignore[import]
|
||
return _get_config()
|
||
except ImportError:
|
||
try:
|
||
from utils.config_loader import get_config as _get_config
|
||
return _get_config()
|
||
except ImportError:
|
||
# Only fallback if truly standalone
|
||
import toml
|
||
class FallbackConfig:
|
||
def __init__(self):
|
||
self.config = {}
|
||
def get(self, key, default=None):
|
||
return default
|
||
def get_section(self, section):
|
||
return {}
|
||
return FallbackConfig()
|
||
|
||
|
||
class QueryAugmenter:
|
||
"""
|
||
Generate augmented queries using various techniques including API-based translation
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
augmentation_methods: Optional[List[str]] = None,
|
||
device: Optional[str] = None,
|
||
config_path: Optional[str] = None
|
||
):
|
||
"""
|
||
Initialize QueryAugmenter.
|
||
Args:
|
||
augmentation_methods: List of augmentation methods to use.
|
||
device: Device to use for augmentation models.
|
||
config_path: Optional path to config file.
|
||
"""
|
||
try:
|
||
self.config = get_config()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load config: {e}, using minimal defaults")
|
||
self.config = get_config() # Always use get_config()
|
||
# Load configuration
|
||
aug_config = self.config.get_section('augmentation')
|
||
api_config = self.config.get_section('api')
|
||
# Set augmentation methods with fallback
|
||
if augmentation_methods is None:
|
||
self.augmentation_methods = aug_config.get('augmentation_methods', ['paraphrase', 'back_translation'])
|
||
else:
|
||
self.augmentation_methods = augmentation_methods
|
||
# Validate augmentation methods (extensible)
|
||
valid_methods = set(['paraphrase', 'back_translation', 'synonym', 'template'])
|
||
self.augmentation_methods = [m for m in self.augmentation_methods if m in valid_methods or callable(getattr(self, f'_augment_{m}', None))]
|
||
if not self.augmentation_methods:
|
||
logger.warning("No valid augmentation methods specified, using default")
|
||
self.augmentation_methods = ['paraphrase']
|
||
# Set device with fallback
|
||
if device is None:
|
||
try:
|
||
self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
|
||
except Exception:
|
||
self.device = 'cpu'
|
||
else:
|
||
self.device = device
|
||
# Configuration parameters with safe defaults
|
||
self.api_config = api_config
|
||
self.augmentation_temperature = aug_config.get('augmentation_temperature', 0.8)
|
||
self.num_augmentations_per_sample = aug_config.get('num_augmentations_per_sample', 2)
|
||
self.enable_api = aug_config.get('enable_query_augmentation', False)
|
||
self.api_first = aug_config.get('api_first', True)
|
||
# API key check (env var override)
|
||
api_key = os.environ.get('BGE_API_KEY') or self.api_config.get('api_key', None)
|
||
if self.enable_api:
|
||
if not api_key or api_key == 'your-api-key-here':
|
||
raise ValueError("API-based augmentation is enabled, but a valid API key is missing in config.toml [api] section or BGE_API_KEY env var.")
|
||
# Model cache
|
||
self.models = {}
|
||
self.paraphrase_model = None
|
||
self.paraphrase_tokenizer = None
|
||
# Load models for augmentation
|
||
self._load_models()
|
||
|
||
def _load_models(self):
|
||
"""Load required models for augmentation with comprehensive error handling"""
|
||
logger.info(f"Loading augmentation models for methods: {self.augmentation_methods}")
|
||
# Load paraphrase model if needed
|
||
if 'paraphrase' in self.augmentation_methods:
|
||
config = get_config()
|
||
model_source = config.get('model_paths.source', 'huggingface')
|
||
try:
|
||
if model_source == 'huggingface':
|
||
logger.info("Loading T5 model for paraphrasing from HuggingFace...")
|
||
self.paraphrase_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
|
||
self.paraphrase_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
|
||
elif model_source == 'modelscope':
|
||
try:
|
||
from modelscope.models import Model as MSModel # type: ignore[import]
|
||
from modelscope.hub.snapshot_download import snapshot_download as ms_snapshot_download # type: ignore[import]
|
||
ms_snapshot_download(model_id="damo/nlp_t5_paraphrase_chinese-base", cache_dir=None)
|
||
self.paraphrase_model = MSModel.from_pretrained("damo/nlp_t5_paraphrase_chinese-base")
|
||
self.paraphrase_tokenizer = None # Add ModelScope tokenizer if needed
|
||
logger.info("T5 model loaded from ModelScope.")
|
||
except ImportError:
|
||
logger.error("modelscope not available. Install it first: pip install modelscope")
|
||
self.paraphrase_model = None
|
||
self.paraphrase_tokenizer = None
|
||
if self.enable_api is False:
|
||
raise RuntimeError("ModelScope paraphrase model required but not available and API augmentation is not enabled.")
|
||
return
|
||
else:
|
||
logger.error(
|
||
f"Unknown model source: {model_source}. Supported: 'huggingface', 'modelscope'.")
|
||
self.paraphrase_model = None
|
||
self.paraphrase_tokenizer = None
|
||
if not self.enable_api:
|
||
logger.warning("Paraphrase model unavailable and API augmentation disabled")
|
||
return
|
||
try:
|
||
if self.paraphrase_model is not None:
|
||
self.paraphrase_model.to(self.device)
|
||
self.paraphrase_model.eval()
|
||
logger.info(f"T5 model loaded successfully on {self.device}")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to move T5 model to {self.device}: {e}, using CPU")
|
||
self.device = 'cpu'
|
||
if self.paraphrase_model is not None:
|
||
self.paraphrase_model.to(self.device)
|
||
self.paraphrase_model.eval()
|
||
except Exception as e:
|
||
logger.error(f"Failed to load T5 model: {e}.")
|
||
self.paraphrase_model = None
|
||
self.paraphrase_tokenizer = None
|
||
if not self.enable_api:
|
||
logger.warning("Paraphrase model failed to load and API augmentation disabled")
|
||
# SOTA back-translation model loading (if available)
|
||
if 'back_translation' in self.augmentation_methods:
|
||
# Example: Use Helsinki-NLP/opus-mt-zh-en and en-zh for Chinese-English back-translation
|
||
try:
|
||
self.bt_zh2en_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
|
||
self.bt_zh2en_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-zh-en")
|
||
self.bt_en2zh_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
|
||
self.bt_en2zh_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-zh")
|
||
logger.info("Loaded SOTA back-translation models for zh<->en.")
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load SOTA back-translation models: {e}")
|
||
self.bt_zh2en_model = None
|
||
self.bt_en2zh_model = None
|
||
|
||
def augment_queries(
|
||
self,
|
||
data: List[Dict],
|
||
num_augmentations: Optional[int] = None,
|
||
keep_original: bool = True
|
||
) -> List[Dict]:
|
||
"""
|
||
Augment queries in the dataset with comprehensive error handling.
|
||
Args:
|
||
data: Original data
|
||
num_augmentations: Number of augmentations per query
|
||
keep_original: Whether to keep original queries
|
||
Returns:
|
||
Augmented dataset
|
||
"""
|
||
if not data:
|
||
logger.warning("Empty data provided for augmentation")
|
||
return []
|
||
|
||
if num_augmentations is None:
|
||
num_augmentations = self.num_augmentations_per_sample
|
||
|
||
if not isinstance(num_augmentations, int) or num_augmentations <= 0:
|
||
logger.warning("Invalid num_augmentations, setting to 1")
|
||
num_augmentations = 1
|
||
|
||
augmented_data = []
|
||
|
||
try:
|
||
for item in tqdm(data, desc="Augmenting queries"):
|
||
# Validate item
|
||
if not isinstance(item, dict) or 'query' not in item:
|
||
logger.warning(f"Invalid item format: {item}")
|
||
continue
|
||
|
||
# Keep original
|
||
if keep_original:
|
||
augmented_data.append(item)
|
||
|
||
# Generate augmentations
|
||
for i in range(num_augmentations):
|
||
try:
|
||
method = random.choice(self.augmentation_methods)
|
||
aug_item = self._augment_single(item, method)
|
||
if aug_item:
|
||
augmented_data.append(aug_item)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to augment item with {method}: {e}")
|
||
continue
|
||
|
||
except Exception as e:
|
||
logger.error(f"Query augmentation failed: {e}")
|
||
# Return original data on complete failure
|
||
return data
|
||
|
||
logger.info(f"Augmented {len(data)} items to {len(augmented_data)} items")
|
||
return augmented_data
|
||
|
||
def _augment_single(self, item: Dict, method: str) -> Optional[Dict]:
|
||
"""Augment a single item using specified method with error handling."""
|
||
if not item.get('query'):
|
||
return None
|
||
|
||
try:
|
||
aug_item = item.copy()
|
||
original_query = item['query']
|
||
|
||
if method == 'paraphrase':
|
||
aug_query = self._paraphrase_query(original_query)
|
||
elif method == 'back_translation':
|
||
aug_query = self._back_translate_query(original_query)
|
||
elif method == 'synonym':
|
||
aug_query = self._synonym_replacement(original_query)
|
||
elif method == 'template':
|
||
aug_query = self._template_based_augmentation(original_query)
|
||
else:
|
||
logger.warning(f"Unknown augmentation method: {method}")
|
||
return None
|
||
|
||
if aug_query and aug_query != original_query:
|
||
aug_item['query'] = aug_query
|
||
aug_item['augmentation_method'] = method
|
||
aug_item['is_augmented'] = True
|
||
return aug_item
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Error in {method} augmentation: {e}")
|
||
|
||
return None
|
||
|
||
def _get_prompt(self, task: str, query: str, lang: str) -> str:
|
||
"""Get prompt for LLM-based augmentation."""
|
||
if lang == 'zh':
|
||
if task == 'paraphrase':
|
||
return f"请将下列文本改写为意思相同的句子，只输出改写后的内容，不要解释：\n\n{query}"
|
||
elif task == 'synonym':
|
||
return f"请将下列文本中的词语替换为合适的同义词，只输出结果，不要解释：\n\n{query}"
|
||
elif task == 'template':
|
||
return f"请用不同的提问模板改写下列问题，只输出结果，不要解释：\n\n{query}"
|
||
else:
|
||
if task == 'paraphrase':
|
||
return f"Paraphrase the following text. Only provide the paraphrase, no explanations:\n\n{query}"
|
||
elif task == 'synonym':
|
||
return f"Replace words in the following text with appropriate synonyms. Only provide the result, no explanations:\n\n{query}"
|
||
elif task == 'template':
|
||
return f"Rephrase the following text using a different question template. Only provide the result, no explanations:\n\n{query}"
|
||
return query
|
||
|
||
def _paraphrase_query(self, query: str) -> str:
|
||
"""Paraphrase a query using the loaded model or API."""
|
||
if not query or not query.strip():
|
||
return query
|
||
lang = 'zh' if self._is_chinese(query) else 'en'
|
||
# Try LLM API first if enabled and configured
|
||
if self.enable_api and self.api_first:
|
||
try:
|
||
prompt = self._get_prompt('paraphrase', query, lang)
|
||
result = self._call_llm_api(prompt)
|
||
if result and result != query:
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"LLM API paraphrasing failed: {e}. Trying local model.")
|
||
# Try T5-based paraphrasing next (only for English)
|
||
if lang == 'en' and self.paraphrase_model is not None and self.paraphrase_tokenizer is not None:
|
||
try:
|
||
input_text = f"paraphrase: {query}"
|
||
inputs = self.paraphrase_tokenizer(
|
||
input_text,
|
||
return_tensors="pt",
|
||
max_length=512,
|
||
truncation=True
|
||
).to(self.device)
|
||
with torch.no_grad():
|
||
outputs = self.paraphrase_model.generate(
|
||
**inputs,
|
||
max_length=150,
|
||
num_return_sequences=1,
|
||
temperature=self.augmentation_temperature,
|
||
do_sample=True,
|
||
pad_token_id=self.paraphrase_tokenizer.eos_token_id
|
||
)
|
||
paraphrase = self.paraphrase_tokenizer.decode(outputs[0], skip_special_tokens=True)
|
||
if paraphrase and paraphrase != query:
|
||
return paraphrase
|
||
except Exception as e:
|
||
logger.warning(f"T5 paraphrasing failed: {e}. Using rule-based approach.")
|
||
# Fallback to rule-based paraphrasing
|
||
return self._rule_based_paraphrase(query, lang)
|
||
|
||
def _rule_based_paraphrase(self, query: str, lang: str = 'en') -> str:
|
||
"""Rule-based paraphrasing for fallback."""
|
||
if not query:
|
||
return query
|
||
try:
|
||
if lang == 'zh':
|
||
patterns = [
|
||
("什么是", "请解释"), ("如何", "怎样"), ("为什么", "什么原因"), ("在哪里", "什么地方"),
|
||
("是什么", "指的是什么"), ("有什么", "存在哪些"), ("能否", "可以"), ("应该", "需要"),
|
||
("必须", "一定要"), ("关于", "有关")
|
||
]
|
||
paraphrased = query
|
||
for original, replacement in patterns:
|
||
if original in query:
|
||
paraphrased = query.replace(original, replacement, 1)
|
||
break
|
||
if paraphrased == query:
|
||
prefixes = ["请问", "我想知道", "能否告诉我", "请教一下", "想了解"]
|
||
paraphrased = f"{random.choice(prefixes)}{query}"
|
||
return paraphrased
|
||
else:
|
||
templates = [
|
||
"What is meant by {}?", "Can you explain {}?", "Tell me about {}", "I want to know about {}",
|
||
"Please describe {}", "Could you clarify {}?", "What does {} refer to?", "How would you define {}?"
|
||
]
|
||
core = query.strip("?.,!").lower()
|
||
question_starters = ["what", "how", "why", "when", "where", "who", "which"]
|
||
for starter in question_starters:
|
||
if core.startswith(starter):
|
||
core = core[len(starter):].strip()
|
||
for aux in ["is", "are", "was", "were", "does", "do", "did"]:
|
||
if core.startswith(aux):
|
||
core = core[len(aux):].strip()
|
||
break
|
||
break
|
||
if core:
|
||
return random.choice(templates).format(core)
|
||
else:
|
||
return query
|
||
except Exception as e:
|
||
logger.warning(f"Rule-based paraphrasing failed: {e}")
|
||
return query
|
||
|
||
def _back_translate_query(self, query: str) -> str:
|
||
"""Back-translate a query using SOTA models if available, else fallback to API or simulation. Always returns a str."""
|
||
if hasattr(self, 'bt_zh2en_model') and self.bt_zh2en_model is not None and self.bt_en2zh_model is not None:
|
||
try:
|
||
is_chinese = any('\u4e00' <= c <= '\u9fff' for c in query)
|
||
if is_chinese:
|
||
inputs = self.bt_zh2en_tokenizer(query, return_tensors="pt", truncation=True, max_length=64).to(self.device)
|
||
en = self.bt_zh2en_model.generate(**inputs, max_length=64)
|
||
en_text = self.bt_zh2en_tokenizer.decode(en[0], skip_special_tokens=True)
|
||
inputs = self.bt_en2zh_tokenizer(en_text, return_tensors="pt", truncation=True, max_length=64).to(self.device)
|
||
zh = self.bt_en2zh_model.generate(**inputs, max_length=64)
|
||
zh_text = self.bt_en2zh_tokenizer.decode(zh[0], skip_special_tokens=True)
|
||
return zh_text if zh_text else query
|
||
else:
|
||
inputs = self.bt_en2zh_tokenizer(query, return_tensors="pt", truncation=True, max_length=64).to(self.device)
|
||
zh = self.bt_en2zh_model.generate(**inputs, max_length=64)
|
||
zh_text = self.bt_en2zh_tokenizer.decode(zh[0], skip_special_tokens=True)
|
||
inputs = self.bt_zh2en_tokenizer(zh_text, return_tensors="pt", truncation=True, max_length=64).to(self.device)
|
||
en = self.bt_zh2en_model.generate(**inputs, max_length=64)
|
||
en_text = self.bt_zh2en_tokenizer.decode(en[0], skip_special_tokens=True)
|
||
return en_text if en_text else query
|
||
except Exception as e:
|
||
logger.warning(f"SOTA back-translation failed: {e}")
|
||
if self.enable_api:
|
||
try:
|
||
is_chinese = any('\u4e00' <= c <= '\u9fff' for c in query)
|
||
result = self._translate_with_api(query, 'zh' if is_chinese else 'en', 'en' if is_chinese else 'zh')
|
||
if result:
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"API back-translation failed: {e}")
|
||
fallback = self._simulate_back_translation(query)
|
||
return fallback if fallback is not None else query
|
||
|
||
def _translate_with_api(self, text: str, source_lang: str, target_lang: str) -> Optional[str]:
|
||
"""Translate text using API."""
|
||
if not text or not text.strip():
|
||
return None
|
||
|
||
base_url = self.api_config.get('base_url', '')
|
||
api_key = self.api_config.get('api_key', '')
|
||
model = self.api_config.get('model', 'deepseek-chat')
|
||
timeout = self.api_config.get('timeout', 30)
|
||
max_retries = self.api_config.get('max_retries', 3)
|
||
|
||
if not base_url or not api_key or api_key == 'your-api-key-here':
|
||
return None
|
||
|
||
headers = {
|
||
'Authorization': f'Bearer {api_key}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
|
||
# Create translation prompt
|
||
lang_names = {'zh': 'Chinese', 'en': 'English', 'ja': 'Japanese', 'fr': 'French'}
|
||
source_name = lang_names.get(source_lang, source_lang)
|
||
target_name = lang_names.get(target_lang, target_lang)
|
||
|
||
prompt = f"Translate the following {source_name} text to {target_name}. Only provide the translation, no explanations:\n\n{text}"
|
||
|
||
data = {
|
||
'model': model,
|
||
'messages': [
|
||
{'role': 'user', 'content': prompt}
|
||
],
|
||
'temperature': self.augmentation_temperature,
|
||
'max_tokens': 500
|
||
}
|
||
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = requests.post(
|
||
f"{base_url}/chat/completions",
|
||
headers=headers,
|
||
json=data,
|
||
timeout=timeout
|
||
)
|
||
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
if 'choices' in result and result['choices']:
|
||
translated = result['choices'][0]['message']['content'].strip()
|
||
return translated
|
||
else:
|
||
logger.warning(f"Invalid response format: {result}")
|
||
else:
|
||
logger.warning(f"Translation API error: {response.status_code} - {response.text}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(2 ** attempt) # Exponential backoff
|
||
|
||
except requests.exceptions.Timeout:
|
||
logger.warning(f"Translation API timeout on attempt {attempt + 1}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(2 ** attempt)
|
||
except requests.exceptions.RequestException as e:
|
||
logger.warning(f"Translation API request failed: {e}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(2 ** attempt)
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error in translation API: {e}")
|
||
break
|
||
|
||
return None
|
||
|
||
def _simulate_back_translation(self, query: str) -> str:
|
||
"""Simulate back-translation as a fallback."""
|
||
if not query:
|
||
return query
|
||
|
||
try:
|
||
# Add language-specific variations that simulate translation effects
|
||
if self._is_chinese(query):
|
||
# Simulate en->zh->en effects
|
||
substitutions = [
|
||
("什么", "哪些"), ("如何", "怎么样"), ("为什么", "什么原因"),
|
||
("非常", "很"), ("特别", "特殊"), ("重要", "关键"),
|
||
("问题", "事情"), ("方法", "方式"), ("内容", "信息")
|
||
]
|
||
else:
|
||
# Simulate zh->en->zh effects
|
||
substitutions = [
|
||
("what", "which"), ("how", "in what way"), ("very", "extremely"),
|
||
("important", "crucial"), ("method", "approach"), ("content", "information"),
|
||
("question", "issue"), ("problem", "matter"), ("find", "locate")
|
||
]
|
||
|
||
result = query
|
||
for original, replacement in substitutions:
|
||
if original in result.lower():
|
||
# Case-preserving replacement
|
||
if original.isupper():
|
||
replacement = replacement.upper()
|
||
elif original.istitle():
|
||
replacement = replacement.title()
|
||
|
||
result = result.replace(original, replacement, 1)
|
||
break
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Back-translation simulation failed: {e}")
|
||
return query
|
||
|
||
def _synonym_replacement(self, query: str) -> str:
|
||
"""Replace words in query with synonyms."""
|
||
if not query or not query.strip():
|
||
return query
|
||
lang = 'zh' if self._is_chinese(query) else 'en'
|
||
# Try LLM API first if enabled and configured
|
||
if self.enable_api and self.api_first:
|
||
try:
|
||
prompt = self._get_prompt('synonym', query, lang)
|
||
result = self._call_llm_api(prompt)
|
||
if result and result != query:
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"LLM API synonym replacement failed: {e}. Using rule-based approach.")
|
||
# Fallback to hard-coded synonym replacement
|
||
try:
|
||
if lang == 'zh':
|
||
synonyms = {
|
||
"重要": ["关键", "核心", "主要"], "方法": ["办法", "手段", "途径"], "问题": ["疑问", "难题", "困惑"],
|
||
"内容": ["材料", "信息", "资料"], "增加": ["提升", "增长", "扩展"], "发现": ["查找", "寻找", "探索"]
|
||
}
|
||
words = list(query)
|
||
replaced = False
|
||
for i, word in enumerate(words):
|
||
if word in synonyms:
|
||
words[i] = random.choice(synonyms[word])
|
||
replaced = True
|
||
if replaced:
|
||
return ''.join(words)
|
||
return query
|
||
else:
|
||
synonyms = {
|
||
"important": ["crucial", "vital", "significant"], "method": ["approach", "technique", "procedure"],
|
||
"find": ["locate", "discover", "identify"], "problem": ["issue", "challenge", "difficulty"],
|
||
"increase": ["rise", "growth", "expansion"], "content": ["material", "substance", "information"]
|
||
}
|
||
words = query.split()
|
||
replaced = False
|
||
for i, word in enumerate(words):
|
||
key = word.lower().strip('.,?!')
|
||
if key in synonyms:
|
||
words[i] = random.choice(synonyms[key])
|
||
replaced = True
|
||
if replaced:
|
||
return ' '.join(words)
|
||
return query
|
||
except Exception as e:
|
||
logger.warning(f"Rule-based synonym replacement failed: {e}")
|
||
return query
|
||
|
||
def _template_based_augmentation(self, query: str) -> str:
|
||
"""Template-based augmentation for queries."""
|
||
if not query or not query.strip():
|
||
return query
|
||
lang = 'zh' if self._is_chinese(query) else 'en'
|
||
# Try LLM API first if enabled and configured
|
||
if self.enable_api and self.api_first:
|
||
try:
|
||
prompt = self._get_prompt('template', query, lang)
|
||
result = self._call_llm_api(prompt)
|
||
if result and result != query:
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"LLM API template-based augmentation failed: {e}. Using rule-based approach.")
|
||
# Fallback to hard-coded templates
|
||
try:
|
||
if lang == 'zh':
|
||
templates = [
|
||
"请找到关于{}的段落", "搜索有关{}的信息", "寻找讨论{}的内容", "我需要了解{}的详细信息", "显示与{}相关的文本"
|
||
]
|
||
core = query.strip("?.,!。：")
|
||
for prefix in ["什么是", "如何", "为什么", "在哪里", "谁是"]:
|
||
if core.startswith(prefix):
|
||
core = core[len(prefix):].strip()
|
||
break
|
||
if core:
|
||
return random.choice(templates).format(core)
|
||
else:
|
||
return query
|
||
else:
|
||
templates = [
|
||
"What is meant by {}?", "Can you explain {}?", "Tell me about {}", "I want to know about {}",
|
||
"Please describe {}", "Could you clarify {}?", "What does {} refer to?", "How would you define {}?"
|
||
]
|
||
core = query.strip("?.,!").lower()
|
||
if core:
|
||
return random.choice(templates).format(core)
|
||
else:
|
||
return query
|
||
except Exception as e:
|
||
logger.warning(f"Rule-based template augmentation failed: {e}")
|
||
return query
|
||
|
||
def _call_llm_api(self, prompt: str) -> Optional[str]:
|
||
"""Call LLM API for augmentation."""
|
||
base_url = self.api_config.get('base_url', '')
|
||
api_key = self.api_config.get('api_key', '')
|
||
model = self.api_config.get('model', 'deepseek-chat')
|
||
timeout = self.api_config.get('timeout', 30)
|
||
max_retries = self.api_config.get('max_retries', 3)
|
||
temperature = self.api_config.get('temperature', self.augmentation_temperature)
|
||
if not base_url or not api_key or api_key == 'your-api-key-here':
|
||
return None
|
||
headers = {
|
||
'Authorization': f'Bearer {api_key}',
|
||
'Content-Type': 'application/json'
|
||
}
|
||
data = {
|
||
'model': model,
|
||
'messages': [
|
||
{'role': 'user', 'content': prompt}
|
||
],
|
||
'temperature': temperature,
|
||
'max_tokens': 500
|
||
}
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = requests.post(
|
||
f"{base_url}/chat/completions",
|
||
headers=headers,
|
||
json=data,
|
||
timeout=timeout
|
||
)
|
||
if response.status_code == 200:
|
||
result = response.json()
|
||
if 'choices' in result and result['choices']:
|
||
content = result['choices'][0]['message']['content'].strip()
|
||
return content
|
||
else:
|
||
logger.warning(f"Invalid response format: {result}")
|
||
else:
|
||
logger.warning(f"LLM API error: {response.status_code} - {response.text}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(2 ** attempt)
|
||
except requests.exceptions.Timeout:
|
||
logger.warning(f"LLM API timeout on attempt {attempt + 1}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(2 ** attempt)
|
||
except requests.exceptions.RequestException as e:
|
||
logger.warning(f"LLM API request failed: {e}")
|
||
if attempt < max_retries - 1:
|
||
time.sleep(2 ** attempt)
|
||
except Exception as e:
|
||
logger.error(f"Unexpected error in LLM API: {e}")
|
||
break
|
||
return None
|
||
|
||
def _is_chinese(self, text: str) -> bool:
|
||
"""Detect if text is Chinese."""
|
||
if not text:
|
||
return False
|
||
try:
|
||
return bool(re.search(r'[\u4e00-\u9fff]', text))
|
||
except Exception:
|
||
return False
|
||
|
||
|
||
class PassageAugmenter:
|
||
"""
|
||
Augment passages to create training variations
|
||
"""
|
||
|
||
def __init__(self):
|
||
try:
|
||
self.config = get_config()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load config: {e}, using defaults")
|
||
self.config = get_config() # Always use get_config()
|
||
|
||
aug_config = self.config.get_section('augmentation')
|
||
self.enable_passage_augmentation = aug_config.get('enable_passage_augmentation', False)
|
||
|
||
self.augmentation_strategies = [
|
||
'sentence_shuffle',
|
||
'span_masking',
|
||
'paraphrase_sentences',
|
||
'add_noise'
|
||
]
|
||
|
||
def augment_passages(
|
||
self,
|
||
data: List[Dict],
|
||
augmentation_prob: Optional[float] = None
|
||
) -> List[Dict]:
|
||
"""
|
||
Augment passages in the dataset with comprehensive error handling
|
||
|
||
Args:
|
||
data: Original data
|
||
augmentation_prob: Probability of augmenting each item
|
||
|
||
Returns:
|
||
Augmented dataset
|
||
"""
|
||
if not self.enable_passage_augmentation:
|
||
logger.info("Passage augmentation disabled in config")
|
||
return data
|
||
|
||
if not data:
|
||
logger.warning("Empty data provided for passage augmentation")
|
||
return []
|
||
|
||
if augmentation_prob is None:
|
||
augmentation_prob = 0.3
|
||
|
||
if not (0 <= augmentation_prob <= 1):
|
||
logger.warning("Invalid augmentation_prob, setting to 0.3")
|
||
augmentation_prob = 0.3
|
||
|
||
augmented_data = []
|
||
|
||
try:
|
||
for item in tqdm(data, desc="Augmenting passages"):
|
||
# Validate item
|
||
if not isinstance(item, dict):
|
||
logger.warning(f"Invalid item format: {item}")
|
||
continue
|
||
|
||
augmented_data.append(item)
|
||
|
||
if random.random() < augmentation_prob:
|
||
# Augment positive passages
|
||
if 'pos' in item and item['pos']:
|
||
try:
|
||
aug_item = item.copy()
|
||
aug_item['pos'] = [
|
||
self._augment_passage(p, random.choice(self.augmentation_strategies))
|
||
for p in item['pos']
|
||
if p and p.strip() # Skip empty passages
|
||
]
|
||
if aug_item['pos']: # Only add if we have valid augmented passages
|
||
aug_item['is_augmented'] = True
|
||
aug_item['augmentation_type'] = 'passage'
|
||
augmented_data.append(aug_item)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to augment passage for item: {e}")
|
||
continue
|
||
|
||
except Exception as e:
|
||
logger.error(f"Passage augmentation failed: {e}")
|
||
return data
|
||
|
||
return augmented_data
|
||
|
||
def _augment_passage(self, passage: str, strategy: str) -> str:
|
||
"""Augment a single passage with error handling"""
|
||
if not passage or not passage.strip():
|
||
return passage
|
||
|
||
try:
|
||
if strategy == 'sentence_shuffle':
|
||
return self._shuffle_sentences(passage)
|
||
elif strategy == 'span_masking':
|
||
return self._mask_spans(passage)
|
||
elif strategy == 'paraphrase_sentences':
|
||
return self._paraphrase_sentences(passage)
|
||
elif strategy == 'add_noise':
|
||
return self._add_noise(passage)
|
||
else:
|
||
logger.warning(f"Unknown augmentation strategy: {strategy}")
|
||
return passage
|
||
except Exception as e:
|
||
logger.warning(f"Error in passage augmentation with {strategy}: {e}")
|
||
return passage
|
||
|
||
def _shuffle_sentences(self, passage: str) -> str:
|
||
"""Shuffle sentences in passage"""
|
||
if not passage or len(passage.strip()) < 10:
|
||
return passage
|
||
|
||
try:
|
||
# Split by Chinese or English sentence endings
|
||
if re.search(r'[\u4e00-\u9fff]', passage):
|
||
sentences = re.split(r'[。！？]', passage)
|
||
delimiter = '。'
|
||
else:
|
||
sentences = re.split(r'[.!?]', passage)
|
||
delimiter = '. '
|
||
|
||
sentences = [s.strip() for s in sentences if s.strip()]
|
||
|
||
if len(sentences) > 2:
|
||
# Keep first and last, shuffle middle
|
||
middle = sentences[1:-1]
|
||
random.shuffle(middle)
|
||
sentences = [sentences[0]] + middle + [sentences[-1]]
|
||
|
||
# Rejoin with appropriate punctuation
|
||
if delimiter == '。':
|
||
return delimiter.join(sentences) + '。'
|
||
else:
|
||
return delimiter.join(sentences) + '.'
|
||
|
||
return passage
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Sentence shuffling failed: {e}")
|
||
return passage
|
||
|
||
def _mask_spans(self, passage: str, mask_ratio: float = 0.15) -> str:
|
||
"""Mask random spans in passage"""
|
||
if not passage or len(passage.strip()) < 20:
|
||
return passage
|
||
|
||
try:
|
||
# Split into words/characters
|
||
if re.search(r'[\u4e00-\u9fff]', passage):
|
||
# Chinese: split by characters for finer granularity
|
||
units = list(passage)
|
||
join_char = ''
|
||
else:
|
||
# English: split by words
|
||
units = passage.split()
|
||
join_char = ' '
|
||
|
||
num_masks = max(1, int(len(units) * mask_ratio))
|
||
|
||
for _ in range(num_masks):
|
||
if len(units) > 5: # Ensure minimum length
|
||
start = random.randint(0, len(units) - 3)
|
||
length = random.randint(1, min(3, len(units) - start))
|
||
for i in range(start, start + length):
|
||
if i < len(units):
|
||
units[i] = '[MASK]'
|
||
|
||
# Rejoin
|
||
return join_char.join(units)
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Span masking failed: {e}")
|
||
return passage
|
||
|
||
def _paraphrase_sentences(self, passage: str) -> str:
|
||
"""Paraphrase some sentences"""
|
||
if not passage or len(passage.strip()) < 20:
|
||
return passage
|
||
|
||
try:
|
||
# Split sentences
|
||
if re.search(r'[\u4e00-\u9fff]', passage):
|
||
sentences = re.split(r'[。！？]', passage)
|
||
delimiter = '。'
|
||
else:
|
||
sentences = re.split(r'[.!?]', passage)
|
||
delimiter = '. '
|
||
|
||
sentences = [s.strip() for s in sentences if s.strip()]
|
||
|
||
if sentences:
|
||
# Randomly modify one sentence
|
||
idx = random.randint(0, len(sentences) - 1)
|
||
original = sentences[idx]
|
||
|
||
# Simple paraphrasing by adding modifiers
|
||
if re.search(r'[\u4e00-\u9fff]', original):
|
||
modifiers = ["显然", "据了解", "实际上", "通常情况下", "根据研究"]
|
||
sentences[idx] = f"{random.choice(modifiers)}，{original}"
|
||
else:
|
||
modifiers = ["Generally", "According to studies", "Typically", "In fact", "Usually"]
|
||
sentences[idx] = f"{random.choice(modifiers)}, {original.lower()}"
|
||
|
||
result = delimiter.join(sentences)
|
||
if delimiter == '。':
|
||
result += '。'
|
||
else:
|
||
result += '.'
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Sentence paraphrasing failed: {e}")
|
||
return passage
|
||
|
||
def _add_noise(self, passage: str, noise_level: float = 0.02) -> str:
|
||
"""Add character-level noise"""
|
||
if not passage or len(passage.strip()) < 10:
|
||
return passage
|
||
|
||
try:
|
||
chars = list(passage)
|
||
num_noisy = max(1, int(len(chars) * noise_level))
|
||
|
||
if len(chars) > 10: # Only add noise to sufficiently long passages
|
||
noise_positions = random.sample(range(len(chars)), min(num_noisy, len(chars) // 5))
|
||
|
||
for pos in noise_positions:
|
||
if pos < len(chars):
|
||
noise_type = random.choice(['duplicate', 'replace'])
|
||
|
||
if noise_type == 'duplicate':
|
||
chars[pos] = chars[pos] * 2
|
||
elif noise_type == 'replace':
|
||
# Replace with similar character
|
||
if re.search(r'[\u4e00-\u9fff]', chars[pos]):
|
||
# Chinese characters
|
||
similar_chars = ['的', '了', '在', '是', '我', '有', '他', '这', '中', '大']
|
||
chars[pos] = random.choice(similar_chars)
|
||
else:
|
||
# English characters
|
||
if chars[pos].isalpha():
|
||
chars[pos] = random.choice('abcdefghijklmnopqrstuvwxyz')
|
||
|
||
return ''.join(chars)
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Noise addition failed: {e}")
|
||
return passage
|
||
|
||
|
||
class HardNegativeAugmenter:
|
||
"""
|
||
Create hard negatives through various strategies
|
||
"""
|
||
|
||
def __init__(self):
|
||
try:
|
||
self.config = get_config()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load config: {e}, using defaults")
|
||
self.config = get_config() # Always use get_config()
|
||
|
||
mining_config = self.config.get_section('hard_negative_mining')
|
||
self.enable_mining = mining_config.get('enable_mining', True)
|
||
self.num_hard_negatives = mining_config.get('num_hard_negatives', 15)
|
||
|
||
self.strategies = [
|
||
'entity_swap',
|
||
'number_change',
|
||
'negation_addition',
|
||
'similar_context'
|
||
]
|
||
|
||
def create_hard_negatives(
|
||
self,
|
||
data: List[Dict],
|
||
num_hard_negatives: Optional[int] = None
|
||
) -> List[Dict]:
|
||
"""
|
||
Create hard negatives for training data with comprehensive error handling
|
||
|
||
Args:
|
||
data: Original data
|
||
num_hard_negatives: Number of hard negatives to create
|
||
|
||
Returns:
|
||
Data with hard negatives
|
||
"""
|
||
if not self.enable_mining:
|
||
logger.info("Hard negative mining disabled in config")
|
||
return data
|
||
|
||
if not data:
|
||
logger.warning("Empty data provided for hard negative creation")
|
||
return []
|
||
|
||
if num_hard_negatives is None:
|
||
num_hard_negatives = self.num_hard_negatives
|
||
|
||
if not isinstance(num_hard_negatives, int) or num_hard_negatives <= 0:
|
||
logger.warning("Invalid num_hard_negatives, setting to 5")
|
||
num_hard_negatives = 5
|
||
|
||
augmented_data = []
|
||
|
||
try:
|
||
for item in tqdm(data, desc="Creating hard negatives"):
|
||
# Validate item
|
||
if not isinstance(item, dict):
|
||
logger.warning(f"Invalid item format: {item}")
|
||
continue
|
||
|
||
aug_item = item.copy()
|
||
|
||
if 'pos' in item and item['pos']:
|
||
hard_negs = []
|
||
|
||
try:
|
||
for pos_passage in item['pos'][:2]: # Use first 2 positives
|
||
if not pos_passage or not pos_passage.strip():
|
||
continue
|
||
|
||
for i in range(min(num_hard_negatives, len(self.strategies) * 2)):
|
||
strategy = self.strategies[i % len(self.strategies)]
|
||
hard_neg = self._create_hard_negative(pos_passage, strategy)
|
||
if hard_neg and hard_neg != pos_passage and hard_neg.strip():
|
||
hard_negs.append(hard_neg)
|
||
|
||
if hard_negs:
|
||
if 'hard_neg' not in aug_item:
|
||
aug_item['hard_neg'] = []
|
||
# Remove duplicates and limit
|
||
unique_negs = list(dict.fromkeys(hard_negs))
|
||
aug_item['hard_neg'].extend(unique_negs[:num_hard_negatives])
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Failed to create hard negatives for item: {e}")
|
||
|
||
augmented_data.append(aug_item)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Hard negative creation failed: {e}")
|
||
return data
|
||
|
||
return augmented_data
|
||
|
||
def _create_hard_negative(self, passage: str, strategy: str) -> str:
|
||
"""Create a hard negative using specified strategy with error handling"""
|
||
if not passage or not passage.strip():
|
||
return passage
|
||
|
||
try:
|
||
if strategy == 'entity_swap':
|
||
return self._swap_entities(passage)
|
||
elif strategy == 'number_change':
|
||
return self._change_numbers(passage)
|
||
elif strategy == 'negation_addition':
|
||
return self._add_negation(passage)
|
||
elif strategy == 'similar_context':
|
||
return self._create_similar_context(passage)
|
||
else:
|
||
logger.warning(f"Unknown hard negative strategy: {strategy}")
|
||
return passage
|
||
except Exception as e:
|
||
logger.warning(f"Error creating hard negative with {strategy}: {e}")
|
||
return passage
|
||
|
||
def _swap_entities(self, passage: str) -> str:
|
||
"""Swap entities to create factually incorrect passage"""
|
||
if not passage:
|
||
return passage
|
||
|
||
try:
|
||
if re.search(r'[\u4e00-\u9fff]', passage):
|
||
entity_pairs = [
|
||
("北京", "上海"), ("中国", "日本"), ("2023年", "2024年"),
|
||
("第一", "第二"), ("增加", "减少"), ("成功", "失败"),
|
||
("重要", "次要"), ("大型", "小型"), ("新的", "旧的")
|
||
]
|
||
else:
|
||
entity_pairs = [
|
||
("Beijing", "Shanghai"), ("China", "Japan"), ("2023", "2024"),
|
||
("first", "second"), ("increase", "decrease"), ("success", "failure"),
|
||
("important", "minor"), ("large", "small"), ("new", "old")
|
||
]
|
||
|
||
aug_passage = passage
|
||
for ent1, ent2 in entity_pairs:
|
||
if ent1 in passage and ent2 not in passage:
|
||
aug_passage = aug_passage.replace(ent1, ent2, 1)
|
||
break
|
||
|
||
return aug_passage
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Entity swapping failed: {e}")
|
||
return passage
|
||
|
||
def _change_numbers(self, passage: str) -> str:
|
||
"""Change numbers in passage"""
|
||
if not passage:
|
||
return passage
|
||
|
||
try:
|
||
def replace_number(match):
|
||
try:
|
||
num = int(match.group())
|
||
# Change by 20-80%
|
||
change = random.uniform(0.2, 0.8)
|
||
if random.random() > 0.5:
|
||
new_num = int(num * (1 + change))
|
||
else:
|
||
new_num = int(num * (1 - change))
|
||
return str(max(1, new_num)) # Ensure positive
|
||
except:
|
||
return match.group()
|
||
|
||
# Find and replace numbers
|
||
aug_passage = re.sub(r'\b\d+\b', replace_number, passage)
|
||
return aug_passage
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Number changing failed: {e}")
|
||
return passage
|
||
|
||
def _add_negation(self, passage: str) -> str:
|
||
"""Add negation to change meaning"""
|
||
if not passage:
|
||
return passage
|
||
|
||
try:
|
||
if re.search(r'[\u4e00-\u9fff]', passage):
|
||
negation_patterns = [
|
||
("是", "不是"), ("有", "没有"), ("会", "不会"),
|
||
("能", "不能"), ("可以", "不可以"), ("应该", "不应该"),
|
||
("成功", "失败"), ("正确", "错误"), ("增长", "下降")
|
||
]
|
||
else:
|
||
negation_patterns = [
|
||
("is", "is not"), ("was", "was not"), ("can", "cannot"),
|
||
("will", "will not"), ("should", "should not"), ("has", "has no"),
|
||
("successful", "unsuccessful"), ("correct", "incorrect"), ("increase", "decrease")
|
||
]
|
||
|
||
aug_passage = passage
|
||
for positive, negative in negation_patterns:
|
||
if positive in passage and negative not in passage:
|
||
aug_passage = aug_passage.replace(positive, negative, 1)
|
||
break
|
||
|
||
return aug_passage
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Negation addition failed: {e}")
|
||
return passage
|
||
|
||
def _create_similar_context(self, passage: str) -> str:
|
||
"""Create passage with similar context but different meaning"""
|
||
if not passage:
|
||
return passage
|
||
|
||
try:
|
||
# Split into sentences
|
||
if re.search(r'[\u4e00-\u9fff]', passage):
|
||
sentences = re.split(r'[。！？]', passage)
|
||
delimiter = '。'
|
||
generic_statements = [
|
||
"这一点在许多情况下都有所体现",
|
||
"相关研究表明了不同的结果",
|
||
"实际情况可能有所不同",
|
||
"这个观点存在一定的争议",
|
||
"具体细节仍在讨论之中"
|
||
]
|
||
else:
|
||
sentences = re.split(r'[.!?]', passage)
|
||
delimiter = '. '
|
||
generic_statements = [
|
||
"This point is reflected in many cases",
|
||
"Related research shows different results",
|
||
"The actual situation may vary",
|
||
"This viewpoint is somewhat controversial",
|
||
"Specific details are still under discussion"
|
||
]
|
||
|
||
sentences = [s.strip() for s in sentences if s.strip()]
|
||
|
||
if len(sentences) >= 3:
|
||
# Replace middle sentence with generic statement
|
||
middle_idx = len(sentences) // 2
|
||
sentences[middle_idx] = random.choice(generic_statements)
|
||
|
||
result = delimiter.join(sentences)
|
||
if delimiter == '。':
|
||
result += '。'
|
||
else:
|
||
result += '.'
|
||
|
||
return result
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Similar context creation failed: {e}")
|
||
return passage
|
||
|
||
|
||
class CrossLingualAugmenter:
|
||
"""
|
||
Create cross-lingual training data with enhanced error handling
|
||
"""
|
||
|
||
def __init__(self,
|
||
source_lang: str = 'zh',
|
||
target_langs: Optional[List[str]] = None):
|
||
try:
|
||
self.config = get_config()
|
||
except Exception as e:
|
||
logger.warning(f"Failed to load config: {e}, using defaults")
|
||
self.config = get_config() # Always use get_config()
|
||
|
||
self.source_lang = source_lang
|
||
self.target_langs = target_langs or ['en']
|
||
|
||
# Get API configuration
|
||
self.api_config = self.config.get_section('api')
|
||
self.enable_api = bool(self.api_config.get('api_key') and
|
||
self.api_config.get('api_key') != 'your-api-key-here')
|
||
|
||
def create_cross_lingual_pairs(
|
||
self,
|
||
data: List[Dict],
|
||
translation_prob: float = 0.3
|
||
) -> List[Dict]:
|
||
"""
|
||
Create cross-lingual query-passage pairs with comprehensive error handling
|
||
|
||
Args:
|
||
data: Original data
|
||
translation_prob: Probability of creating translation
|
||
|
||
Returns:
|
||
Augmented data with cross-lingual pairs
|
||
"""
|
||
if not data:
|
||
logger.warning("Empty data provided for cross-lingual augmentation")
|
||
return []
|
||
|
||
if not (0 <= translation_prob <= 1):
|
||
logger.warning("Invalid translation_prob, setting to 0.3")
|
||
translation_prob = 0.3
|
||
|
||
augmented_data = []
|
||
|
||
try:
|
||
for item in tqdm(data, desc="Creating cross-lingual pairs"):
|
||
# Validate item
|
||
if not isinstance(item, dict) or 'query' not in item:
|
||
logger.warning(f"Invalid item format: {item}")
|
||
continue
|
||
|
||
augmented_data.append(item)
|
||
|
||
if random.random() < translation_prob:
|
||
# Create translated query version
|
||
for target_lang in self.target_langs:
|
||
try:
|
||
trans_item = item.copy()
|
||
trans_query = self._translate_query(
|
||
item['query'],
|
||
self.source_lang,
|
||
target_lang
|
||
)
|
||
if trans_query and trans_query != item['query']:
|
||
trans_item['query'] = trans_query
|
||
trans_item['cross_lingual'] = True
|
||
trans_item['query_lang'] = target_lang
|
||
trans_item['passage_lang'] = self.source_lang
|
||
augmented_data.append(trans_item)
|
||
except Exception as e:
|
||
logger.warning(f"Failed to create cross-lingual pair for {target_lang}: {e}")
|
||
continue
|
||
|
||
except Exception as e:
|
||
logger.error(f"Cross-lingual augmentation failed: {e}")
|
||
return data
|
||
|
||
return augmented_data
|
||
|
||
def _translate_query(self, query: str, source: str, target: str) -> str:
|
||
"""Translate query from source to target language with error handling"""
|
||
if not query or not query.strip():
|
||
return query
|
||
|
||
# Try API translation first if enabled
|
||
if self.enable_api:
|
||
try:
|
||
# Use the existing API translation method
|
||
augmenter = QueryAugmenter()
|
||
result = augmenter._translate_with_api(query, source, target)
|
||
if result:
|
||
return result
|
||
except Exception as e:
|
||
logger.warning(f"API translation failed: {e}")
|
||
|
||
# Fallback to rule-based simulation
|
||
return self._simulate_translation(query, source, target)
|
||
|
||
def _simulate_translation(self, query: str, source: str, target: str) -> str:
|
||
"""Simulate translation with pattern-based replacements"""
|
||
if not query:
|
||
return query
|
||
|
||
try:
|
||
if source == 'zh' and target == 'en':
|
||
# Chinese to English patterns
|
||
replacements = {
|
||
"什么": "what", "如何": "how", "为什么": "why", "哪里": "where",
|
||
"谁": "who", "书": "book", "章节": "chapter", "信息": "information",
|
||
"内容": "content", "文本": "text", "段落": "passage"
|
||
}
|
||
translated = query
|
||
for zh, en in replacements.items():
|
||
translated = translated.replace(zh, en)
|
||
return f"[{source}→{target}] {translated}"
|
||
|
||
elif source == 'en' and target == 'zh':
|
||
# English to Chinese patterns
|
||
replacements = {
|
||
"what": "什么", "how": "如何", "why": "为什么", "where": "哪里",
|
||
"who": "谁", "book": "书", "chapter": "章节", "information": "信息",
|
||
"content": "内容", "text": "文本", "passage": "段落"
|
||
}
|
||
translated = query.lower()
|
||
for en, zh in replacements.items():
|
||
translated = translated.replace(en, zh)
|
||
return f"[{source}→{target}] {translated}"
|
||
|
||
# For other language pairs, just add marker
|
||
return f"[{source}→{target}] {query}"
|
||
|
||
except Exception as e:
|
||
logger.warning(f"Translation simulation failed: {e}")
|
||
return query
|