2024-06-11 07:06:29 +02:00
package main
import (
"flag"
"fmt"
2024-06-11 07:29:31 +02:00
"os"
2024-06-12 08:01:44 +02:00
"strings"
2024-06-11 07:06:29 +02:00
)
var (
2024-06-13 09:01:55 +02:00
inputFile = flag . String ( "i" , "" , "input file (markdown, or plain text)" )
model = flag . String ( "m" , "llama3" , "llm model" )
outputFile = flag . String ( "o" , "" , "output file" )
2024-06-11 07:06:29 +02:00
)
2024-06-12 08:01:44 +02:00
const (
promptStart = ` Could you translate the following text into English? Try te preserve the original tone of voice as much as possible. Only answer with the translation itself, no additional comments needed. `
)
2024-06-11 07:06:29 +02:00
func main ( ) {
flag . Parse ( )
2024-06-11 07:29:31 +02:00
if * inputFile == "" {
fmt . Println ( "No input file specified" )
os . Exit ( 1 )
}
ollamaHost := os . Getenv ( "OLLAMA_HOST" )
if ollamaHost == "" {
fmt . Println ( "OLLAMA_HOST environment variable not set" )
os . Exit ( 1 )
}
ollama := NewOllama ( ollamaHost )
2024-06-12 08:01:44 +02:00
doc , err := os . ReadFile ( * inputFile )
2024-06-11 07:29:31 +02:00
if err != nil {
fmt . Println ( err )
os . Exit ( 1 )
}
2024-06-11 07:06:29 +02:00
2024-06-13 09:01:55 +02:00
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 ) )
2024-06-12 08:01:44 +02:00
translated := make ( [ ] string , 0 )
2024-06-13 09:01:55 +02:00
for _ , chunk := range chunks {
2024-06-12 08:01:44 +02:00
prompt := fmt . Sprintf ( "%s\n---\n%s" , promptStart , chunk )
res , err := ollama . Generate ( * model , prompt )
if err != nil {
fmt . Println ( err )
os . Exit ( 1 )
}
fmt . Printf ( "." )
translated = append ( translated , res )
}
2024-06-13 09:01:55 +02:00
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 )
}
2024-06-12 08:01:44 +02:00
fmt . Printf ( "\n\n%s\n" , strings . Join ( translated , "\n" ) )
2024-06-11 07:06:29 +02:00
}