Jay Brown
106 lines
2.3 KiB
Go
106 lines
2.3 KiB
Go
//go:build aws
|
|
|
|
//go:generate go run -tags=aws ./textdetection.go
|
|
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"io"
|
|
"log/slog"
|
|
"os"
|
|
"time"
|
|
|
|
documenttext "queryorchestration/internal/document/text"
|
|
documenttypes "queryorchestration/internal/document/types"
|
|
"queryorchestration/internal/serviceconfig"
|
|
"queryorchestration/internal/serviceconfig/aws"
|
|
"queryorchestration/internal/serviceconfig/objectstore"
|
|
"queryorchestration/internal/serviceconfig/textract"
|
|
)
|
|
|
|
// Config provides the configuration for the text generator test tool.
|
|
// Note: Query-related configuration has been removed. See remove_query_plan.md.
|
|
type Config struct {
|
|
serviceconfig.BaseConfig
|
|
textract.TextractConfig
|
|
objectstore.ObjectStoreConfig
|
|
}
|
|
|
|
func main() {
|
|
ctx := context.Background()
|
|
|
|
if len(os.Args) != 2 {
|
|
fmt.Println("Error: expected one argument")
|
|
fmt.Println("Usage:", os.Args[0], "<file_base>")
|
|
os.Exit(1)
|
|
}
|
|
|
|
name := os.Args[1]
|
|
|
|
slog.Info("arguments", "file_base", name)
|
|
|
|
cfg := &Config{}
|
|
profile := aws.Profile(os.Getenv("AWS_PROFILE"))
|
|
if profile == "" {
|
|
slog.Error("profile not available")
|
|
os.Exit(1)
|
|
}
|
|
|
|
err := cfg.SetTextractClientWithProfile(ctx, profile)
|
|
if err != nil {
|
|
slog.Error("error creating client", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
root, err := os.OpenRoot("assets")
|
|
if err != nil {
|
|
slog.Error("error opening root", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
pdfFile, err := root.Open(fmt.Sprintf("original/%s.pdf", name))
|
|
if err != nil {
|
|
slog.Error("error opening pdf file", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
pdf, err := documenttypes.NewPDFFromReader(pdfFile)
|
|
if err != nil {
|
|
slog.Error("error creating pdf file", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
textSvc := documenttext.New(cfg)
|
|
|
|
startTime := time.Now()
|
|
|
|
text, err := textSvc.GetDocumentText(ctx, pdf)
|
|
if err != nil {
|
|
slog.Error("error getting page count", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
outputFilename := fmt.Sprintf("text/%s.txt", name)
|
|
file, err := root.Create(outputFilename)
|
|
if err != nil {
|
|
slog.Error("error creating file", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
defer file.Close()
|
|
|
|
s, err := io.ReadAll(text)
|
|
if err != nil {
|
|
slog.Error("error reading text", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
_, err = file.Write(s)
|
|
if err != nil {
|
|
slog.Error("error writing file", "error", err)
|
|
os.Exit(1)
|
|
}
|
|
|
|
slog.Info("extracted text", "output", outputFilename, "time", time.Since(startTime))
|
|
}
|