diff --git a/ollama-translate/.gitignore b/ollama-translate/.gitignore index 7545a50..2e1fa2d 100644 --- a/ollama-translate/.gitignore +++ b/ollama-translate/.gitignore @@ -1 +1 @@ -test.md \ No newline at end of file +*.md \ No newline at end of file diff --git a/ollama-translate/main.go b/ollama-translate/main.go index 86ed095..7de008e 100644 --- a/ollama-translate/main.go +++ b/ollama-translate/main.go @@ -8,8 +8,9 @@ import ( ) var ( - inputFile = flag.String("i", "", "input file") - model = flag.String("m", "llama3", "model file") + inputFile = flag.String("i", "", "input file (markdown, or plain text)") + model = flag.String("m", "llama3", "llm model") + outputFile = flag.String("o", "", "output file") ) const ( @@ -35,8 +36,22 @@ func main() { os.Exit(1) } + chunks := make([]string, 0) + for _, par := range strings.Split(string(doc), "\n\n") { + last := len(chunks) - 1 + switch { + case last == -1: + chunks = append(chunks, par) + case last >= 0 && len(par)+len(chunks[last]) > 500: + chunks = append(chunks, par) + default: + chunks[last] = fmt.Sprintf("%s\n\n%s", chunks[last], par) + } + } + + fmt.Printf("translating %d chunks\n", len(chunks)) translated := make([]string, 0) - for _, chunk := range strings.Split(string(doc), "\n\n") { + for _, chunk := range chunks { prompt := fmt.Sprintf("%s\n---\n%s", promptStart, chunk) res, err := ollama.Generate(*model, prompt) if err != nil { @@ -47,5 +62,14 @@ func main() { translated = append(translated, res) } + if *outputFile != "" { + if err := os.WriteFile(*outputFile, []byte(strings.Join(translated, "\n")), 0644); err != nil { + fmt.Println(err) + os.Exit(1) + } + fmt.Printf("\nfile written to %s\n", *outputFile) + os.Exit(0) + } + fmt.Printf("\n\n%s\n", strings.Join(translated, "\n")) }