303 lines
8.8 KiB
Go
303 lines
8.8 KiB
Go
package main
|
|
|
|
import (
|
|
"flag"
|
|
"fmt"
|
|
"os"
|
|
"path/filepath"
|
|
"strings"
|
|
"sync"
|
|
|
|
"github.com/schollz/progressbar/v3"
|
|
"github.com/wholetrans/d10n/config"
|
|
"github.com/wholetrans/d10n/translator"
|
|
)
|
|
|
|
func main() {
|
|
// Parse command line arguments
|
|
if len(os.Args) < 2 {
|
|
printUsage()
|
|
os.Exit(1)
|
|
}
|
|
|
|
sourcePath := os.Args[1]
|
|
|
|
// Set up flags
|
|
flagSet := flag.NewFlagSet("d10n", flag.ExitOnError)
|
|
|
|
targetPathPtr := flagSet.String("target", "", "Target path for translated content")
|
|
sourceLanguagePtr := flagSet.String("source-lang", "", "Source language code")
|
|
targetLanguagePtr := flagSet.String("target-lang", "", "Target language code (required)")
|
|
modelPtr := flagSet.String("model", "", "Model to use for translation")
|
|
apiKeyPtr := flagSet.String("api-key", "", "API key for OpenAI compatible service")
|
|
apiBasePtr := flagSet.String("api-base", "", "API base URL for OpenAI compatible service")
|
|
systemPromptPtr := flagSet.String("system-prompt", "", "System prompt for the model")
|
|
formatPtr := flagSet.String("format", "", "File format to process (e.g., md, txt)")
|
|
|
|
// Concurrency flag
|
|
concurrencyPtr := flagSet.Int("concurrency", 0, "Number of concurrent translation tasks (default is 3)")
|
|
|
|
// Chunking flags
|
|
chunkEnabledPtr := flagSet.Bool("chunk", false, "Enable chunked translation")
|
|
chunkSizePtr := flagSet.Int("chunk-size", 0, "Size of each chunk in tokens (default is 10240)")
|
|
chunkPromptPtr := flagSet.String("chunk-prompt", "", "Prompt to use for continuing translation (default is 'Please continue translation')")
|
|
chunkContextPtr := flagSet.Int("chunk-context", 0, "Number of chunks to include as context (default is 2)")
|
|
|
|
// Parse flags
|
|
if err := flagSet.Parse(os.Args[2:]); err != nil {
|
|
fmt.Println("Error parsing arguments:", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
// Load configuration
|
|
cfg, err := config.LoadConfig()
|
|
if err != nil {
|
|
fmt.Println("Error loading configuration:", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
// Validate and set defaults
|
|
if *targetLanguagePtr == "" {
|
|
fmt.Println("Error: Target language is required")
|
|
printUsage()
|
|
os.Exit(1)
|
|
}
|
|
|
|
if *sourceLanguagePtr == "" {
|
|
fmt.Println("Error: Source language is required")
|
|
printUsage()
|
|
os.Exit(1)
|
|
}
|
|
|
|
// Override config with command line arguments if provided
|
|
if *modelPtr != "" {
|
|
cfg.Model = *modelPtr
|
|
}
|
|
|
|
if *apiKeyPtr != "" {
|
|
cfg.APIKey = *apiKeyPtr
|
|
}
|
|
|
|
if *apiBasePtr != "" {
|
|
cfg.APIBase = *apiBasePtr
|
|
}
|
|
|
|
if *systemPromptPtr != "" {
|
|
cfg.SystemPrompt = *systemPromptPtr
|
|
}
|
|
|
|
// Set concurrency options
|
|
if *concurrencyPtr > 0 {
|
|
cfg.Concurrency = *concurrencyPtr
|
|
}
|
|
|
|
// Set chunking options from command line arguments
|
|
if *chunkEnabledPtr {
|
|
cfg.Chunk.Enabled = true
|
|
}
|
|
|
|
if *chunkSizePtr > 0 {
|
|
cfg.Chunk.Size = *chunkSizePtr
|
|
}
|
|
|
|
if *chunkPromptPtr != "" {
|
|
cfg.Chunk.Prompt = *chunkPromptPtr
|
|
}
|
|
|
|
if *chunkContextPtr > 0 {
|
|
cfg.Chunk.Context = *chunkContextPtr
|
|
}
|
|
|
|
// Set target path if not provided
|
|
targetPath := *targetPathPtr
|
|
if targetPath == "" {
|
|
ext := filepath.Ext(sourcePath)
|
|
base := strings.TrimSuffix(sourcePath, ext)
|
|
targetPath = base + "_l10n" + ext
|
|
}
|
|
|
|
// Create translator
|
|
trans := translator.NewTranslator(cfg)
|
|
|
|
// Process the source path
|
|
fileInfo, err := os.Stat(sourcePath)
|
|
if err != nil {
|
|
fmt.Printf("Error accessing source path %s: %v\n", sourcePath, err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
if fileInfo.IsDir() {
|
|
err = processDirectory(sourcePath, targetPath, *sourceLanguagePtr, *targetLanguagePtr, *formatPtr, trans)
|
|
} else {
|
|
err = processFile(sourcePath, targetPath, *sourceLanguagePtr, *targetLanguagePtr, trans)
|
|
}
|
|
|
|
if err != nil {
|
|
fmt.Println("Error during processing:", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
fmt.Println("Translation completed successfully!")
|
|
}
|
|
|
|
func printUsage() {
|
|
fmt.Println("Usage: d10n <source_path> [options]")
|
|
fmt.Println("Options:")
|
|
fmt.Println(" -target <path> Target path for translated content")
|
|
fmt.Println(" -source-lang <code> Source language code (required)")
|
|
fmt.Println(" -target-lang <code> Target language code (required)")
|
|
fmt.Println(" -model <model> Model to use for translation")
|
|
fmt.Println(" -api-key <key> API key for OpenAI compatible service")
|
|
fmt.Println(" -api-base <url> API base URL for OpenAI compatible service")
|
|
fmt.Println(" -system-prompt <prompt> System prompt for the model")
|
|
fmt.Println(" -format <ext> File format to process (e.g., md, txt)")
|
|
fmt.Println(" -concurrency <num> Number of concurrent translation tasks (default: 3)")
|
|
fmt.Println(" -chunk Enable chunked translation")
|
|
fmt.Println(" -chunk-size <tokens> Size of each chunk in tokens (default: 10240)")
|
|
fmt.Println(" -chunk-prompt <prompt> Prompt for continuing translation (default: 'Please continue translation')")
|
|
fmt.Println(" -chunk-context <num> Number of chunks to include as context (default: 2)")
|
|
}
|
|
|
|
func processFile(sourcePath, targetPath, sourceLanguage, targetLanguage string, trans *translator.Translator) error {
|
|
content, err := os.ReadFile(sourcePath)
|
|
if err != nil {
|
|
return fmt.Errorf("error reading file %s: %w", sourcePath, err)
|
|
}
|
|
|
|
translatedContent, err := trans.Translate(string(content), sourceLanguage, targetLanguage, sourcePath)
|
|
if err != nil {
|
|
return fmt.Errorf("error translating file %s: %w", sourcePath, err)
|
|
}
|
|
|
|
// Create target directory if it doesn't exist
|
|
targetDir := filepath.Dir(targetPath)
|
|
if err := os.MkdirAll(targetDir, 0755); err != nil {
|
|
return fmt.Errorf("error creating target directory %s: %w", targetDir, err)
|
|
}
|
|
|
|
// Write translated content to target path
|
|
if err := os.WriteFile(targetPath, []byte(translatedContent), 0644); err != nil {
|
|
return fmt.Errorf("error writing to file %s: %w", targetPath, err)
|
|
}
|
|
|
|
// Remove the progress bar for this file
|
|
trans.RemoveProgressBar(sourcePath)
|
|
|
|
return nil
|
|
}
|
|
|
|
func processDirectory(sourcePath, targetPath, sourceLanguage, targetLanguage, format string, trans *translator.Translator) error {
|
|
// Create target directory if it doesn't exist
|
|
if err := os.MkdirAll(targetPath, 0755); err != nil {
|
|
return fmt.Errorf("error creating target directory %s: %w", targetPath, err)
|
|
}
|
|
|
|
// First scan to find all matching files
|
|
var matchedFiles []string
|
|
err := filepath.Walk(sourcePath, func(path string, info os.FileInfo, err error) error {
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Skip directories
|
|
if info.IsDir() {
|
|
return nil
|
|
}
|
|
|
|
// Check if file matches the format
|
|
if format != "" && !strings.HasSuffix(info.Name(), "."+format) {
|
|
return nil
|
|
}
|
|
|
|
matchedFiles = append(matchedFiles, path)
|
|
return nil
|
|
})
|
|
|
|
if err != nil {
|
|
return err
|
|
}
|
|
|
|
// Display all matched files
|
|
fmt.Printf("Found %d files to translate:\n", len(matchedFiles))
|
|
for _, file := range matchedFiles {
|
|
relPath, _ := filepath.Rel(sourcePath, file)
|
|
fmt.Printf("- %s\n", relPath)
|
|
}
|
|
fmt.Println()
|
|
|
|
// Create overall progress bar
|
|
overallBar := progressbar.NewOptions(len(matchedFiles),
|
|
progressbar.OptionSetDescription("[Overall Progress]"),
|
|
progressbar.OptionShowCount(),
|
|
progressbar.OptionSetTheme(progressbar.Theme{
|
|
Saucer: "#",
|
|
SaucerHead: ">",
|
|
SaucerPadding: "-",
|
|
BarStart: "[",
|
|
BarEnd: "]",
|
|
}),
|
|
)
|
|
|
|
// Set up a worker pool for concurrent processing
|
|
var (
|
|
wg sync.WaitGroup
|
|
mutex sync.Mutex
|
|
errors []error
|
|
concurrency = trans.GetConcurrency()
|
|
semaphore = make(chan struct{}, concurrency)
|
|
)
|
|
|
|
// Process files concurrently
|
|
for i, path := range matchedFiles {
|
|
// Compute relative path from source
|
|
relPath, err := filepath.Rel(sourcePath, path)
|
|
if err != nil {
|
|
return fmt.Errorf("error computing relative path for %s: %w", path, err)
|
|
}
|
|
|
|
// Compute target file path
|
|
targetFilePath := filepath.Join(targetPath, relPath)
|
|
|
|
// Create target directory structure if needed
|
|
targetFileDir := filepath.Dir(targetFilePath)
|
|
if err := os.MkdirAll(targetFileDir, 0755); err != nil {
|
|
return fmt.Errorf("error creating target directory %s: %w", targetFileDir, err)
|
|
}
|
|
|
|
// Process individual file concurrently
|
|
wg.Add(1)
|
|
go func(idx int, sourcePath, targetPath, relPath string) {
|
|
defer wg.Done()
|
|
|
|
// Acquire semaphore (limit concurrency)
|
|
semaphore <- struct{}{}
|
|
defer func() { <-semaphore }()
|
|
|
|
// Process the file
|
|
err := processFile(sourcePath, targetPath, sourceLanguage, targetLanguage, trans)
|
|
|
|
// Update overall progress bar
|
|
mutex.Lock()
|
|
overallBar.Add(1)
|
|
if err != nil {
|
|
errors = append(errors, err)
|
|
fmt.Printf("Error translating %s: %v\n", relPath, err)
|
|
}
|
|
mutex.Unlock()
|
|
}(i+1, path, targetFilePath, relPath)
|
|
}
|
|
|
|
// Wait for all translations to complete
|
|
wg.Wait()
|
|
|
|
// Complete the overall progress bar
|
|
overallBar.Finish()
|
|
fmt.Println() // Add some spacing after progress bars
|
|
|
|
// Check if any errors occurred
|
|
if len(errors) > 0 {
|
|
return fmt.Errorf("encountered %d errors during translation", len(errors))
|
|
}
|
|
|
|
return nil
|
|
}
|