Merge pull request #52 from Paulanerus/dev

Dev

Author: Paulanerus

api/src/main/kotlin/dev/paulee/api/data/IDataService.kt

@@ -28,6 +28,10 @@ interface IDataService : Closeable {
 
     suspend fun downloadModel(model: Embedding.Model, path: Path, onProgress: (progress: Int) -> Unit)
 
+    suspend fun importPool(path: Path): Boolean
+
+    suspend fun exportPool(dataInfo: DataInfo, path: Path): Boolean
+
     fun getPage(
         query: String,
         isSemantic: Boolean,

api/src/main/kotlin/dev/paulee/api/internal/Embedding.kt

@@ -15,7 +15,9 @@ object Embedding {
             "Google DeepMind",
             "300M",
             "https://huggingface.co/google/embeddinggemma-300m",
-            ModelData()
+            ModelData(
+                modelData = "onnx/model.onnx_data"
+            )
         ),
         AncientGreekBert(
             "onnx-community/Ancient-Greek-BERT-ONNX",
@@ -25,7 +27,16 @@ object Embedding {
             "https://huggingface.co/pranaydeeps/Ancient-Greek-BERT",
             ModelData(
                 maxLength = 512,
-                modelData = ""
+            )
+        ),
+        GreekTransfer(
+            "onnx-community/stsb-xlm-r-greek-transfer-ONNX",
+            "model_management.greek_transfer.desc",
+            "lighteternal",
+            "270M",
+            "https://huggingface.co/lighteternal/stsb-xlm-r-greek-transfer",
+            ModelData(
+                maxLength = 400,
             )
         )
     }
@@ -34,7 +45,7 @@ object Embedding {
         val dimension: Int = 768,
         val maxLength: Int = 2048,
         val model: String = "onnx/model.onnx",
-        val modelData: String = "onnx/model.onnx_data",
+        val modelData: String? = null,
         val tokenizer: String = "tokenizer.json",
         val tokenizerConfig: String = "tokenizer_config.json",
     )

core/src/main/kotlin/dev/paulee/core/data/DataServiceImpl.kt

@@ -237,6 +237,46 @@ object DataServiceImpl : IDataService {
             val dataInfo = FileService.fromJson(runCatching { jsonFile.readText() }.getOrDefault(""))
                 ?: return@forEachDirectoryEntry
 
+            val dataPath = child.resolve("data")
+
+            if (dataPath.notExists()) {
+                logger.warn("Data pool '${child.name}' has no data directory. Migrating files...")
+
+                runCatching { dataPath.createDirectories() }
+                    .onFailure {
+                        logger.error("Failed to create data directory for ${child.name}.", it)
+                        return@forEachDirectoryEntry
+                    }
+
+                var failed = false
+                dataInfo.sources.forEach { source ->
+                    val name = source.name
+
+                    val sourcePath = FileService.dataDir.resolve("$name.csv")
+
+                    if (sourcePath.notExists()) {
+                        logger.warn("Data pool '${child.name}' has no source file '$name'. Skipping.")
+                        failed = true
+                        return@forEach
+                    }
+
+                    val targetPath = dataPath.resolve("$name.csv")
+
+                    runCatching { sourcePath.moveTo(targetPath, overwrite = true) }
+                        .onFailure {
+                            logger.error("Failed to migrate file '$name'", it)
+                            failed = true
+                        }
+                }
+
+                if (failed) {
+                    logger.error("Data pool '${child.name}' has failed to migrate files. Skipping.")
+                    return@forEachDirectoryEntry
+                }
+
+                logger.info("Migrated data pool '${child.name}' files.")
+            }
+
             val infoName = dataInfo.name
 
             if (dataInfo.storageType == StorageType.SQLITE) {
@@ -255,10 +295,7 @@ object DataServiceImpl : IDataService {
 
                 logger.info("Loaded $infoName data pool.")
             } else {
-                logger.info("Deleting invalid or empty data pool directory '$infoName'.")
-
-                runCatching { child.deleteRecursively() }
-                    .onFailure { e -> logger.error("Failed to delete directory ${child.fileName}.", e) }
+                logger.warn("Data pool '$infoName' has invalid storage provider.")
             }
         }
 
@@ -319,6 +356,187 @@ object DataServiceImpl : IDataService {
     override suspend fun downloadModel(model: Embedding.Model, path: Path, onProgress: (progress: Int) -> Unit) =
         EmbeddingProvider.downloadModel(model, path, onProgress)
 
+    @OptIn(ExperimentalPathApi::class)
+    override suspend fun importPool(path: Path): Boolean {
+        if (path.extension != "zip") {
+            logger.warn("File is not a zip file.")
+            return false
+        }
+
+        logger.info("Importing pool from file '$path'.")
+
+        val pathOut = FileService.tempDir.resolve(path.nameWithoutExtension)
+
+        val resultUnzip = withContext(Dispatchers.IO) { ZipService.unzip(path, pathOut) }
+
+        if (!resultUnzip) {
+            logger.warn("Failed to unzip pool.")
+            return false
+        }
+
+        fun deleteIfExists() {
+            if (pathOut.exists()) pathOut.deleteRecursively()
+        }
+
+        val infoFile = pathOut.resolve("info.json")
+
+        if (infoFile.notExists() || !infoFile.isRegularFile()) {
+            logger.warn("Pool info file does not exist.")
+
+            deleteIfExists()
+
+            return false
+        }
+
+        val dataInfo = dataInfoFromString(infoFile.readText())
+
+        if (dataInfo == null) {
+            logger.warn("Pool info file is invalid.")
+
+            deleteIfExists()
+
+            return false
+        }
+
+        val destinationDir = FileService.dataDir.resolve(dataInfo.name)
+
+        if (destinationDir.exists()) {
+            logger.warn("Pool directory already exists.")
+
+            deleteIfExists()
+
+            return false
+        }
+
+        runCatching { destinationDir.createDirectories() }
+            .onFailure {
+                logger.error("Failed to create pool directory.", it)
+                deleteIfExists()
+
+                return false
+            }
+
+        val indexPath = pathOut.resolve("index")
+
+        if (indexPath.notExists() || !indexPath.isDirectory()) {
+            logger.warn("Pool index directory does not exist.")
+
+            deleteIfExists()
+
+            return false
+        }
+
+        val (resultChecker, report) = Indexer.checkIndex(indexPath)
+
+        if (!resultChecker) {
+            logger.error("Pool index is invalid: $report")
+
+            deleteIfExists()
+
+            return false
+        } else logger.info("Pool index is valid.")
+
+        val storageDir = pathOut.resolve("data")
+
+        if (storageDir.notExists() || !storageDir.isDirectory()) {
+            logger.warn("Pool storage directory does not exist.")
+
+            deleteIfExists()
+
+            return false
+        }
+
+        runCatching { storageDir.moveTo(destinationDir.resolve("data")) }
+            .onFailure {
+                logger.error("Failed to move pool data directory.", it)
+
+                deleteIfExists()
+
+                return false
+            }
+
+        runCatching { infoFile.moveTo(destinationDir.resolve("info.json")) }
+            .onFailure {
+                logger.error("Failed to move pool info file.", it)
+
+                deleteIfExists()
+
+                return false
+            }
+
+        runCatching { indexPath.moveTo(destinationDir.resolve("index")) }
+            .onFailure {
+                logger.error("Failed to move pool index directory.", it)
+
+                deleteIfExists()
+
+                return false
+            }
+
+        logger.info("Initializing storage provider for pool '${dataInfo.name}'...")
+
+        val storageProvider = StorageProvider.of(dataInfo.storageType)
+
+        val status = storageProvider.init(dataInfo, destinationDir)
+
+        if (status != ProviderStatus.Success) {
+            logger.warn("Failed to initialize storage provider for pool '${dataInfo.name}'.")
+            return false
+        }
+
+        logger.info("Initializing pool...")
+
+        dataPools[dataInfo.name] =
+            runCatching { DataPool(Indexer(destinationDir.resolve("index"), dataInfo), dataInfo, storageProvider) }
+                .getOrElse {
+                    logger.error("Failed to initialize pool.", it)
+                    return false
+                }
+
+        logger.info("Successfully imported pool '${dataInfo.name}'.")
+
+        deleteIfExists()
+
+        return true
+    }
+
+    @OptIn(ExperimentalPathApi::class)
+    override suspend fun exportPool(dataInfo: DataInfo, path: Path): Boolean {
+        logger.info("Exporting pool '${dataInfo.name}' to '$path'.")
+
+        val infoName = dataInfo.name
+
+        val poolPath = FileService.dataDir.resolve(infoName)
+
+        if (poolPath.notExists()) {
+            logger.warn("Pool path does not exist.")
+            return false
+        }
+
+        logger.info("Packaging pool files for '$infoName'.")
+
+        val pathOut = path.resolve("$infoName.zip")
+
+        val result = withContext(Dispatchers.IO) {
+            val resultZip = ZipService.zip(
+                setOf(
+                    poolPath.resolve("index"),
+                    poolPath.resolve("data"),
+                    poolPath.resolve("info.json")
+                ),
+                pathOut,
+                filesToExclude = setOf("write.lock")
+            )
+
+            resultZip
+        }
+
+        if (result) logger.info("Exported pool '$infoName' to '$pathOut'.")
+        else logger.error("Failed to export pool.")
+
+        return result
+    }
+
     override fun getPage(query: String, isSemantic: Boolean, order: QueryOrder?, pageCount: Int): PageResult {
 
         if (this.currentPool == null || this.currentField == null) return Pair(emptyList(), emptyMap())

core/src/main/kotlin/dev/paulee/core/data/FileService.kt

@@ -86,6 +86,8 @@ internal object FileService {
 
     val modelsDir: Path get() = ensureDir("models")
 
+    val tempDir: Path get() = ensureDir("temp")
+
     private val logger = getLogger(FileService::class.java)
 
     private val mapper = jacksonObjectMapper().apply { enable(SerializationFeature.INDENT_OUTPUT) }