diff --git a/.idea/libraries/jetbrains_ml_models_python_imports_ranking_model.xml b/.idea/libraries/jetbrains_ml_models_python_imports_ranking_model.xml deleted file mode 100644 index 3ec5732e6946..000000000000 --- a/.idea/libraries/jetbrains_ml_models_python_imports_ranking_model.xml +++ /dev/null @@ -1,18 +0,0 @@ - - - - - - 0516e6383101eb995b8f4d9c7dc217578bbc7467a806895d5c4e2390c27bba7d - - - - - - - - - - - - \ No newline at end of file diff --git a/.idea/libraries/jetbrains_mlapi_catboost_shadow_need_slf4j.xml b/.idea/libraries/jetbrains_mlapi_catboost_shadow_need_slf4j.xml deleted file mode 100644 index 8b8db7176586..000000000000 --- a/.idea/libraries/jetbrains_mlapi_catboost_shadow_need_slf4j.xml +++ /dev/null @@ -1,16 +0,0 @@ - - - - - - 94a06c3fdd4aa638df59c8c783a379ca43a5b933187760e219e09e84f1633d76 - - - - - - - - - - \ No newline at end of file diff --git a/.idea/libraries/jetbrains_mlapi_extension.xml b/.idea/libraries/jetbrains_mlapi_extension.xml new file mode 100644 index 000000000000..d64aab9bd95a --- /dev/null +++ b/.idea/libraries/jetbrains_mlapi_extension.xml @@ -0,0 +1,20 @@ + + + + + + 1d18aa7a0d891cab230ddb67e3d1174830e1472c7373c37296c23a3905040a7d + + + + + + + + + + + + + + \ No newline at end of file diff --git a/.idea/libraries/jetbrains_mlapi_ml_building_blocks.xml b/.idea/libraries/jetbrains_mlapi_ml_building_blocks.xml deleted file mode 100644 index dc73d41aafb3..000000000000 --- a/.idea/libraries/jetbrains_mlapi_ml_building_blocks.xml +++ /dev/null @@ -1,20 +0,0 @@ - - - - - - 65b9f465f00474bcbfff5949458631a92c6a4cf0587f05fac4d0acd65cfae4cf - - - - - - - - - - - - - - \ No newline at end of file diff --git a/.idea/libraries/jetbrains_mlapi_ml_feature_api.xml b/.idea/libraries/jetbrains_mlapi_ml_feature_api.xml deleted file mode 100644 index 3acd427e7a08..000000000000 --- a/.idea/libraries/jetbrains_mlapi_ml_feature_api.xml +++ /dev/null @@ -1,18 +0,0 @@ - - - - - - 8f0300408884fe7ea8b3665f4bc2e7f725cda763806ae92677cf141f4d781193 - - - - - - - - - - - - \ No newline at end of file diff --git a/.idea/libraries/jetbrains_mlapi_usage.xml b/.idea/libraries/jetbrains_mlapi_usage.xml new file mode 100644 index 000000000000..446cf5c4e9c3 --- /dev/null +++ b/.idea/libraries/jetbrains_mlapi_usage.xml @@ -0,0 +1,20 @@ + + + + + + b28271a408b89f4a9e00e4a849543b5e9206333472ee379027d969e5d1b42e6f + + + + + + + + + + + + + + \ No newline at end of file diff --git a/.idea/modules.xml b/.idea/modules.xml index 92104cb53340..87f82b6f3914 100644 --- a/.idea/modules.xml +++ b/.idea/modules.xml @@ -898,7 +898,6 @@ - diff --git a/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityLibraryLicenses.kt b/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityLibraryLicenses.kt index 25cfd32ef326..c8a27df90f63 100644 --- a/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityLibraryLicenses.kt +++ b/platform/build-scripts/src/org/jetbrains/intellij/build/CommunityLibraryLicenses.kt @@ -1323,9 +1323,8 @@ object CommunityLibraryLicenses { jetbrainsLibrary("intellij.remoterobot.remote.fixtures"), jetbrainsLibrary("intellij.remoterobot.robot.server.core"), jetbrainsLibrary("jetbrains.intellij.deps.rwmutex.idea"), - jetbrainsLibrary("jetbrains.ml.models.python.imports.ranking.model"), - jetbrainsLibrary("jetbrains.mlapi.ml.building.blocks"), - jetbrainsLibrary("jetbrains.mlapi.ml.feature.api"), + jetbrainsLibrary("jetbrains.mlapi.extension"), + jetbrainsLibrary("jetbrains.mlapi.usage"), jetbrainsLibrary("jshell-frontend"), jetbrainsLibrary("jvm-native-trusted-roots"), jetbrainsLibrary("kotlin-gradle-plugin-idea"), diff --git a/platform/build-scripts/src/org/jetbrains/intellij/build/python/PythonCommunityPluginModules.kt b/platform/build-scripts/src/org/jetbrains/intellij/build/python/PythonCommunityPluginModules.kt index c6a3b0fe2da9..aa5c70221c63 100644 --- a/platform/build-scripts/src/org/jetbrains/intellij/build/python/PythonCommunityPluginModules.kt +++ b/platform/build-scripts/src/org/jetbrains/intellij/build/python/PythonCommunityPluginModules.kt @@ -30,7 +30,6 @@ object PythonCommunityPluginModules { "intellij.python.pydev", "intellij.python.sdk", "intellij.python.terminal", - "intellij.python.ml.features", ) /** diff --git a/platform/ml-api/intellij.platform.ml.iml b/platform/ml-api/intellij.platform.ml.iml index 577c2cea61b4..4bf4ebf195b9 100644 --- a/platform/ml-api/intellij.platform.ml.iml +++ b/platform/ml-api/intellij.platform.ml.iml @@ -11,6 +11,6 @@ - + \ No newline at end of file diff --git a/platform/ml-impl/intellij.platform.ml.impl.iml b/platform/ml-impl/intellij.platform.ml.impl.iml index 5505980c18d6..9cbf75142a37 100644 --- a/platform/ml-impl/intellij.platform.ml.impl.iml +++ b/platform/ml-impl/intellij.platform.ml.impl.iml @@ -50,7 +50,6 @@ - - + \ No newline at end of file diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/MLModelSelector.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/MLModelSelector.kt deleted file mode 100644 index cf4c4e47a10b..000000000000 --- a/platform/ml-impl/src/com/intellij/platform/ml/impl/MLModelSelector.kt +++ /dev/null @@ -1,134 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.platform.ml.impl - -import com.intellij.openapi.diagnostic.thisLogger -import com.intellij.openapi.util.registry.Registry -import com.jetbrains.ml.building.blocks.model.MLModel -import com.jetbrains.ml.building.blocks.model.MLModelLoader -import com.jetbrains.ml.building.blocks.task.MLTask -import com.jetbrains.ml.features.api.MLUnit -import com.jetbrains.ml.features.api.feature.FeatureDeclaration -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.async -import org.jetbrains.annotations.ApiStatus -import java.nio.file.Path -import java.util.concurrent.atomic.AtomicReference -import kotlin.io.path.exists - - -/** - * A utility class that orchestrates ML models for a machine learning task. - * - * It is aware of a default [defaultLoader] and when [modelPathRegistryKey] - * is not empty, it loads a model from the corresponding path - */ -@ApiStatus.Internal -open class MLModelSelector, P : Any>( - private val defaultLoader: MLModelLoader, - registryModelLoaderFactory: (Path) -> MLModelLoader, - private val modelPathRegistryKey: String, - private val coroutineScope: CoroutineScope, -) { - private val isUsingRegistry by lazy { Registry.`is`(modelPathRegistryKey) } - private val registryLoader = RegistryModelLoader(modelPathRegistryKey, defaultLoader, registryModelLoaderFactory) - private val modelLoader by lazy { - if (isUsingRegistry) registryLoader else defaultLoader - } - private var state: AtomicReference> = AtomicReference(State.NotLoaded) - - /** - * Starts loading the ML model if it has not been loaded yet. - * Call this function before calling [getModel] to make it - * return not null on the first call, solving the problem of - * cold start. - */ - fun prepareModel( - features: Map, List>>, - parameters: Map? = null, - ) { - when (state.get()) { - is State.Loaded -> return - is State.Loading -> return - is State.NotLoaded -> startLoading(features, parameters) - } - } - - fun prepareModel(task: MLTask, parameters: Map? = null) { - prepareModel(task.treeFeaturesFlattened, parameters) - } - - /** - * Returns an instance of an ML model if it has been loaded. Null otherwise. - * The model starts to load during the first - * - * If registry value of [modelPathRegistryKey] is not empty, the model is loaded from the specified path - */ - fun getModel( - features: Map, List>>, - parameters: Map? = null, - ): M? { - return when (val s = state.get()) { - is State.Loaded -> s.model - is State.Loading -> null - is State.NotLoaded -> { - startLoading(features, parameters) - null - } - } - } - - fun getModel( - task: MLTask, - parameters: Map? = null, - ): M? { - return getModel(task.treeFeaturesFlattened, parameters) - } - - - private fun startLoading( - features: Map, List>>, - parameters: Map?, - ) { - if (!state.compareAndSet(State.NotLoaded, State.Loading)) return - coroutineScope.async(Dispatchers.IO) { - try { - val model = modelLoader.loadModel(features, parameters) - assert(state.compareAndSet(State.Loading, State.Loaded(model))) - } catch (e: Throwable) { - thisLogger().error("Unable to load ML model", e) - state.compareAndSet(State.Loading, State.NotLoaded) - } - } - } - - private sealed interface State, in P : Any> { - data object NotLoaded : State, Any> - data object Loading : State, Any> - class Loaded, P : Any>(val model: M) : State - } -} - -private class RegistryModelLoader, P : Any>( - private val registryKey: String, - defaultLoader: MLModelLoader, - private val registryModelLoaderFactory: (Path) -> MLModelLoader, -) : MLModelLoader { - override val predictionClass: Class

= defaultLoader.predictionClass - - override suspend fun loadModel(features: Map, List>>, parameters: Map?): M { - val key = Registry.get(registryKey) - require(key.isRestartRequired()) { - "Registry key `$registryKey` must require restart" - } - val path = key.asString().takeUnless { it.isEmpty() }?.let { Path.of(it) } - requireNotNull(path) { "Registry key `$registryKey` must contain non-empty string" } - require(path.isAbsolute) { "Registry key `$registryKey` must be an absolute path" } - require(path.exists()) { "Registry key `$registryKey` be an existing path" } - val loader = registryModelLoaderFactory(path) - return loader.loadModel(features, parameters) - } -} - -private val MLTask.treeFeaturesFlattened: Map, List>> - get() = treeFeatures.flatMap { it.entries }.associateBy({ it.key }, { it.value }) diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/IJPlatform.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/IJPlatform.kt new file mode 100644 index 000000000000..fd0de7aaebd0 --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/IJPlatform.kt @@ -0,0 +1,90 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools + +import com.intellij.openapi.components.Service +import com.intellij.openapi.diagnostic.Logger +import com.intellij.openapi.diagnostic.debug +import com.intellij.util.application +import com.intellij.util.messages.Topic +import com.jetbrains.ml.platform.MLApiPlatform.ExtensionController +import org.jetbrains.annotations.ApiStatus + + +@ApiStatus.Internal +@Service(Service.Level.APP) +class IJPlatform : com.jetbrains.ml.platform.MLApiPlatform( + featureProviders = EP_NAME_FEATURE_PROVIDER.extensionList, + mlUnitProviders = emptyList(), +) { + override val taskListeners: Map>> + get() = KeyedMessagingProvider.collect(MLTaskListenerTyped.TOPIC) + + override fun addTaskListener(taskId: String, taskListener: com.jetbrains.ml.monitoring.MLTaskListenerTyped<*, *>): com.jetbrains.ml.platform.MLApiPlatform.ExtensionController { + val connection = application.messageBus.connect() + + fun , P : Any> capturingType(taskListenerTyped: com.jetbrains.ml.monitoring.MLTaskListenerTyped) { + connection.subscribe(MLTaskListenerTyped.TOPIC, object : MessageBusMLTaskListenerProvider { + override fun provide(collector: (com.jetbrains.ml.monitoring.MLTaskListenerTyped, String) -> Unit) = collector(taskListenerTyped, taskId) + }) + } + + capturingType(taskListener) + + return ExtensionController { connection.disconnect() } + } + + override val loggingListeners: Map> + get() = KeyedMessagingProvider.collect(MLTaskLoggingListener.TOPIC) + + override fun addLoggingListener(taskId: String, loggingListener: com.jetbrains.ml.monitoring.MLTaskLoggingListener): ExtensionController { + val connection = application.messageBus.connect() + connection.subscribe(MLTaskLoggingListener.TOPIC, object : MessageBusMLTaskLoggingListenerProvider { + override fun provide(collector: (com.jetbrains.ml.monitoring.MLTaskLoggingListener, String) -> Unit) = collector(loggingListener, taskId) + }) + + return ExtensionController { connection.disconnect() } + } + + override val systemLoggerBuilder: com.jetbrains.ml.platform.SystemLoggerBuilder = object : com.jetbrains.ml.platform.SystemLoggerBuilder { + override fun build(clazz: Class<*>): com.jetbrains.ml.platform.SystemLogger { + return IJSystemLogger(Logger.getInstance(clazz)) + } + + override fun build(name: String): com.jetbrains.ml.platform.SystemLogger { + return IJSystemLogger(Logger.getInstance(name)) + } + + private inner class IJSystemLogger(private val baseLogger: Logger) : com.jetbrains.ml.platform.SystemLogger { + override fun info(data: () -> String) = baseLogger.info(data()) + + override fun warn(data: () -> String) = baseLogger.warn(data()) + + override fun debug(data: () -> String) = baseLogger.debug { data() } + + override fun error(e: Throwable) = baseLogger.error(e) + } + } +} + + +@ApiStatus.Internal +sealed interface KeyedMessagingProvider { + fun provide(collector: (T, String) -> Unit) + + companion object { + fun

, T> collect(topic: Topic

): Map> { + val collected = mutableMapOf>() + application.messageBus.syncPublisher(topic).provide { it, keyOfIt -> + @Suppress("UNCHECKED_CAST") + collected.getOrPut(keyOfIt) { mutableListOf() }.add(it as T) + } + return collected + } + } +} + +@ApiStatus.Internal +interface MessageBusMLTaskListenerProvider, P : Any> : KeyedMessagingProvider> + +@ApiStatus.Internal +interface MessageBusMLTaskLoggingListenerProvider : KeyedMessagingProvider diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/extensionPoints.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/extensionPoints.kt new file mode 100644 index 000000000000..7c1e48f8daf9 --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/extensionPoints.kt @@ -0,0 +1,7 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools + +import com.intellij.openapi.extensions.ExtensionPointName +import com.jetbrains.ml.FeatureProvider + +internal val EP_NAME_FEATURE_PROVIDER: ExtensionPointName = ExtensionPointName.create("com.intellij.platform.ml.featureProvider") diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/listeners.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/listeners.kt new file mode 100644 index 000000000000..4f771b9f57ad --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/listeners.kt @@ -0,0 +1,31 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools + +import com.intellij.util.messages.Topic +import org.jetbrains.annotations.ApiStatus + +@ApiStatus.Internal +interface MLTaskListenerTyped, P : Any> : com.jetbrains.ml.monitoring.MLTaskListenerTyped, MessageBusMLTaskListenerProvider { + val task: com.jetbrains.ml.MLTask + + override fun provide(collector: (com.jetbrains.ml.monitoring.MLTaskListenerTyped, String) -> Unit) { + collector(this, task.id) + } + + companion object { + val TOPIC: Topic> = Topic.create>("ml.task", MessageBusMLTaskListenerProvider::class.java) + } +} + +@ApiStatus.Internal +interface MLTaskLoggingListener : com.jetbrains.ml.monitoring.MLTaskLoggingListener, MessageBusMLTaskLoggingListenerProvider { + val task: com.jetbrains.ml.MLTask<*, *> + + override fun provide(collector: (com.jetbrains.ml.monitoring.MLTaskLoggingListener, String) -> Unit) { + collector(this, task.id) + } + + companion object { + val TOPIC = Topic.create("ml.logging", MessageBusMLTaskLoggingListenerProvider::class.java) + } +} diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logging.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logging.kt new file mode 100644 index 000000000000..41d0b8359105 --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logging.kt @@ -0,0 +1,21 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools + +import com.intellij.internal.statistic.eventLog.EventLogGroup +import com.intellij.openapi.Disposable +import com.intellij.openapi.observable.util.whenDisposed +import com.intellij.platform.ml.impl.tools.logs.IntelliJFusEventRegister +import com.jetbrains.ml.MLTask +import com.jetbrains.ml.model.MLModel +import org.jetbrains.annotations.ApiStatus + +@ApiStatus.Internal +fun , P : Any> EventLogGroup.registerMLTaskLogging( + task: MLTask, + parentDisposable: Disposable, + eventPrefix: String = task.id, +) { + val componentRegister = IntelliJFusEventRegister(this) + val listenerController = componentRegister.registerLogging(task, eventPrefix) + parentDisposable.whenDisposed { listenerController.remove() } +} diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/logs/fus/IntelliJFusEventRegister.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/IntelliJFusEventRegister.kt similarity index 84% rename from platform/ml-impl/src/com/intellij/platform/ml/impl/logs/fus/IntelliJFusEventRegister.kt rename to platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/IntelliJFusEventRegister.kt index 2982420cf2a0..9bd0a5955117 100644 --- a/platform/ml-impl/src/com/intellij/platform/ml/impl/logs/fus/IntelliJFusEventRegister.kt +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/IntelliJFusEventRegister.kt @@ -1,16 +1,16 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.platform.ml.impl.logs.fus +package com.intellij.platform.ml.impl.tools.logs import com.intellij.internal.statistic.eventLog.FeatureUsageData import com.intellij.internal.statistic.eventLog.events.PrimitiveEventField import com.intellij.internal.statistic.eventLog.events.VarargEventId import com.intellij.lang.Language import com.intellij.openapi.util.Version -import com.intellij.platform.ml.impl.logs.fus.ConverterObjectDescription.Companion.asIJObjectDescription -import com.intellij.platform.ml.impl.logs.fus.ConverterOfEnum.Companion.toIJConverter -import com.intellij.platform.ml.impl.logs.fus.IJEventPairConverter.Companion.typedBuild -import com.jetbrains.ml.building.blocks.logs.FusEventLogger -import com.jetbrains.ml.building.blocks.logs.FusEventRegister +import com.intellij.platform.ml.impl.tools.logs.ConverterObjectDescription.Companion.asIJObjectDescription +import com.intellij.platform.ml.impl.tools.logs.ConverterOfEnum.Companion.toIJConverter +import com.intellij.platform.ml.impl.tools.logs.IJEventPairConverter.Companion.typedBuild +import com.jetbrains.ml.logs.schema.EventField +import com.jetbrains.ml.logs.schema.EventPair import org.jetbrains.annotations.ApiStatus import com.intellij.internal.statistic.eventLog.EventLogGroup as IJEventLogGroup import com.intellij.internal.statistic.eventLog.events.BooleanEventField as IJBooleanEventField @@ -28,35 +28,35 @@ import com.intellij.internal.statistic.eventLog.events.ObjectEventData as IJObje import com.intellij.internal.statistic.eventLog.events.ObjectEventField as IJObjectEventField import com.intellij.internal.statistic.eventLog.events.ObjectListEventField as IJObjectListEventField import com.intellij.internal.statistic.eventLog.events.StringEventField as IJStringEventField -import com.jetbrains.ml.features.api.logs.BooleanEventField as MLBooleanEventField -import com.jetbrains.ml.features.api.logs.ClassEventField as MLClassEventField -import com.jetbrains.ml.features.api.logs.DoubleEventField as MLDoubleEventField -import com.jetbrains.ml.features.api.logs.EnumEventField as MLEnumEventField -import com.jetbrains.ml.features.api.logs.EventField as MLEventField -import com.jetbrains.ml.features.api.logs.EventPair as MLEventPair -import com.jetbrains.ml.features.api.logs.FloatEventField as MLFloatEventField -import com.jetbrains.ml.features.api.logs.IntEventField as MLIntEventField -import com.jetbrains.ml.features.api.logs.LongEventField as MLLongEventField -import com.jetbrains.ml.features.api.logs.ObjectDescription as MLObjectDescription -import com.jetbrains.ml.features.api.logs.ObjectEventData as MLObjectEventData -import com.jetbrains.ml.features.api.logs.ObjectEventField as MLObjectEventField -import com.jetbrains.ml.features.api.logs.ObjectListEventField as MLObjectListEventField -import com.jetbrains.ml.features.api.logs.StringEventField as MLStringEventField +import com.jetbrains.ml.logs.schema.BooleanEventField as MLBooleanEventField +import com.jetbrains.ml.logs.schema.ClassEventField as MLClassEventField +import com.jetbrains.ml.logs.schema.DoubleEventField as MLDoubleEventField +import com.jetbrains.ml.logs.schema.EnumEventField as MLEnumEventField +import com.jetbrains.ml.logs.schema.EventField as MLEventField +import com.jetbrains.ml.logs.schema.EventPair as MLEventPair +import com.jetbrains.ml.logs.schema.FloatEventField as MLFloatEventField +import com.jetbrains.ml.logs.schema.IntEventField as MLIntEventField +import com.jetbrains.ml.logs.schema.LongEventField as MLLongEventField +import com.jetbrains.ml.logs.schema.ObjectDescription as MLObjectDescription +import com.jetbrains.ml.logs.schema.ObjectEventData as MLObjectEventData +import com.jetbrains.ml.logs.schema.ObjectEventField as MLObjectEventField +import com.jetbrains.ml.logs.schema.ObjectListEventField as MLObjectListEventField +import com.jetbrains.ml.logs.schema.StringEventField as MLStringEventField @ApiStatus.Internal -class IntelliJFusEventRegister(private val baseEventGroup: IJEventLogGroup) : FusEventRegister { +class IntelliJFusEventRegister(private val baseEventGroup: IJEventLogGroup) : com.jetbrains.ml.logs.FusEventRegister { private class Logger( private val varargEventId: VarargEventId, private val objectDescription: ConverterObjectDescription - ) : FusEventLogger { - override suspend fun log(eventPairs: List>) { + ) : com.jetbrains.ml.logs.FusEventLogger { + override fun log(eventPairs: List>) { val ijEventPairs = objectDescription.buildEventPairs(eventPairs) varargEventId.log(*ijEventPairs.toTypedArray()) } } - override fun registerEvent(name: String, eventFields: List>): FusEventLogger { + override fun registerEvent(name: String, eventFields: List>): com.jetbrains.ml.logs.FusEventLogger { val objectDescription = ConverterObjectDescription(MLObjectDescription(eventFields)) val varargEventId = baseEventGroup.registerVarargEvent(name, null, *objectDescription.getFields()) return Logger(varargEventId, objectDescription) @@ -89,6 +89,7 @@ private fun createConverter(mlEventField: MLEventField): IJEventPairConve is VersionEventField -> ConverterOfVersion(mlEventField) as IJEventPairConverter } } + else -> throw IllegalArgumentException( """ Conversion of ${mlEventField.javaClass.simpleName} is not possible. @@ -101,7 +102,7 @@ private fun createConverter(mlEventField: MLEventField): IJEventPairConve private class ConverterOfCustom(mlEventField: IJCustomEventField) : IJEventPairConverter { override val ijEventField: IJEventField = mlEventField.baseIJEventField - override fun buildEventPair(mlEventPair: MLEventPair): IJEventPair { + override fun buildEventPair(mlEventPair: EventPair): IJEventPair { return ijEventField with mlEventPair.data } } @@ -113,7 +114,7 @@ private class ConverterOfString(mlEventField: MLStringEventField) : IJEventPairC description = mlEventField.lazyDescription() ) - override fun buildEventPair(mlEventPair: MLEventPair): IJEventPair { + override fun buildEventPair(mlEventPair: EventPair): IJEventPair { return ijEventField with mlEventPair.data } } @@ -126,7 +127,7 @@ private class ConverterOfLanguage(mlEventField: LanguageEventField) : IJEventPai } } -private class ConverterOfVersion(mlEventField: com.intellij.platform.ml.impl.logs.fus.VersionEventField) : IJEventPairConverter { +private class ConverterOfVersion(mlEventField: com.intellij.platform.ml.impl.tools.logs.VersionEventField) : IJEventPairConverter { private class VersionEventField(override val name: String, override val description: String?) : PrimitiveEventField() { override val validationRule: List get() = listOf("{regexp#version}") diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/LanguageSpecific.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/LanguageSpecific.kt new file mode 100644 index 000000000000..40238772e170 --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/LanguageSpecific.kt @@ -0,0 +1,32 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools.logs + +import com.intellij.lang.Language +import com.jetbrains.ml.logs.schema.EventField +import org.jetbrains.annotations.ApiStatus + +/** + * Something, that is dedicated for one language only. + */ +@ApiStatus.Internal +interface LanguageSpecific { + val language: Language +} + +/** + * The analyzer, that adds information about ML model's language to logs. + */ +@ApiStatus.Internal +class ModelLanguageAnalyser : com.jetbrains.ml.analysis.MLTaskAnalyserTyped + where M : com.jetbrains.ml.model.MLModel

, + M : LanguageSpecific { + private val LANGUAGE = LanguageEventField("model_language") { "The programming language the ML model is trained for" } + + override fun startMLSessionAnalysis(sessionInfo: com.jetbrains.ml.session.MLSessionInfo) = object : com.jetbrains.ml.analysis.MLSessionAnalyserTyped { + override suspend fun analyseSession(tree: com.jetbrains.ml.tree.MLTree.ATopNode?) = buildList { + tree?.mlModel?.let { add(LANGUAGE with it.language) } + } + } + + override val declaration: List> = listOf(LANGUAGE) +} diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/Versioned.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/Versioned.kt new file mode 100644 index 000000000000..703a68042908 --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/Versioned.kt @@ -0,0 +1,34 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools.logs + +import com.intellij.openapi.util.Version +import com.jetbrains.ml.logs.schema.EventField +import org.jetbrains.annotations.ApiStatus + +/** + * Something, that has versions. + */ +@ApiStatus.Internal +interface Versioned { + val version: Version? +} + +/** + * Adds model's version to the ML logs. + */ +@ApiStatus.Internal +class ModelVersionAnalyser : com.jetbrains.ml.analysis.MLTaskAnalyserTyped + where M : com.jetbrains.ml.model.MLModel

, + M : Versioned { + companion object { + private val VERSION = VersionEventField("model_version") { "Version of the ML model" } + } + + override fun startMLSessionAnalysis(sessionInfo: com.jetbrains.ml.session.MLSessionInfo) = object : com.jetbrains.ml.analysis.MLSessionAnalyserTyped { + override suspend fun analyseSession(tree: com.jetbrains.ml.tree.MLTree.ATopNode?) = buildList { + tree?.mlModel?.version?.let { add(VERSION with it) } + } + } + + override val declaration: List> = listOf(VERSION) +} diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/logs/fus/eventFields.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/eventFields.kt similarity index 89% rename from platform/ml-impl/src/com/intellij/platform/ml/impl/logs/fus/eventFields.kt rename to platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/eventFields.kt index 4b0838bed0cd..15021f87bfdd 100644 --- a/platform/ml-impl/src/com/intellij/platform/ml/impl/logs/fus/eventFields.kt +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/logs/eventFields.kt @@ -1,13 +1,12 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.platform.ml.impl.logs.fus +package com.intellij.platform.ml.impl.tools.logs import com.intellij.lang.Language import com.intellij.openapi.util.Version -import com.jetbrains.ml.features.api.logs.CustomRuleEventField +import com.jetbrains.ml.logs.schema.CustomRuleEventField import org.jetbrains.annotations.ApiStatus import com.intellij.internal.statistic.eventLog.events.EventField as IJEventField - @ApiStatus.Internal sealed class IJSpecificEventField(name: String, lazyDescription: () -> String) : CustomRuleEventField(name, lazyDescription) diff --git a/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/model/RegressionModel.kt b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/model/RegressionModel.kt new file mode 100644 index 000000000000..8ed44312757e --- /dev/null +++ b/platform/ml-impl/src/com/intellij/platform/ml/impl/tools/model/RegressionModel.kt @@ -0,0 +1,142 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.platform.ml.impl.tools.model + +import com.intellij.internal.ml.DecisionFunction +import com.jetbrains.ml.* +import com.jetbrains.ml.model.MLModel +import org.jetbrains.annotations.ApiStatus + + +/** + * A wrapper for using legacy [DecisionFunction] ML API`s [MLModel]. + */ +@ApiStatus.Internal +open class RegressionModel private constructor( + private val decisionFunction: DecisionFunction, + private val modelUnits: Set>, + globalUnits: Set>, + private val featureSerialization: FeatureNameSerialization, +) : MLModel { + constructor( + decisionFunction: DecisionFunction, + featureSerialization: FeatureNameSerialization, + globalUnits: Set>, + ) : this( + decisionFunction = decisionFunction, + modelUnits = decisionFunction.featuresOrder.map { featureMapper -> + featureSerialization.deserialize(featureMapper.featureName, globalUnits.associateBy { it.name }).first + }.toSet(), + globalUnits = globalUnits, + featureSerialization = featureSerialization + ) + + override val knownFeatures: Map, FeatureFilter> = createKnownFeatures( + DecisionFunctionWrapper(decisionFunction, globalUnits, featureSerialization), + modelUnits + ) + + override fun predict(taskUnits: List, contexts: List, features: Map, List>): Double { + val array = DoubleArray(decisionFunction.featuresOrder.size) + val featurePerSerializedName: Map = features + .flatMap { (unit, unitFeatures) -> unitFeatures.map { unit to it } } + .associate { (unit, feature) -> featureSerialization.serialize(unit, feature.declaration.name) to feature } + + require(features.keys == modelUnits) { + "Given features units are ${features.keys}, but this model needs ${modelUnits}" + } + + for (featureI in decisionFunction.featuresOrder.indices) { + val featureMapper = decisionFunction.featuresOrder[featureI] + val featureSerializedName = featureMapper.featureName + val featureValue = featureMapper.asArrayValue(featurePerSerializedName[featureSerializedName]?.value) + array[featureI] = featureValue + } + + return decisionFunction.predict(array) + } + + interface FeatureNameSerialization { + fun serialize(mlUnit: MLUnit<*>, featureName: String): String + + fun deserialize(serializedFeatureName: String, unitsPerName: Map>): Pair, String> + } + + object DefaultFeatureSerialization : FeatureNameSerialization { + private const val SERIALIZED_FEATURE_SEPARATOR = '/' + + override fun serialize(mlUnit: MLUnit<*>, featureName: String): String { + return mlUnit.name + SERIALIZED_FEATURE_SEPARATOR + featureName + } + + override fun deserialize(serializedFeatureName: String, unitsPerName: Map>): Pair, String> { + val indexOfLastSeparator = serializedFeatureName.indexOfLast { it == SERIALIZED_FEATURE_SEPARATOR } + require(indexOfLastSeparator >= 0) { "Feature name '$serializedFeatureName' does not contain ML unit's name" } + val unitName = serializedFeatureName.slice(0 until indexOfLastSeparator) + val featureName = serializedFeatureName.slice(indexOfLastSeparator until serializedFeatureName.length) + val featureUnit = requireNotNull(unitsPerName[unitName]) { + """ + Serialized feature '$serializedFeatureName' has tier $unitName, + but all available tiers are ${unitsPerName.keys} + """.trimIndent() + } + return featureUnit to featureName + } + } + + private class DecisionFunctionWrapper( + private val decisionFunction: DecisionFunction, + globalUnits: Set>, + private val featureNameSerialization: FeatureNameSerialization, + ) { + private val availableUnitsPerName: Map> = globalUnits.associateBy { it.name } + + val knownFeatures: Map, Set> = run { + val knownFeaturesSerializedNames = decisionFunction.featuresOrder.map { it.featureName }.toSet() + knownFeaturesSerializedNames + .map { featureNameSerialization.deserialize(it, availableUnitsPerName) } + .groupBy({ it.first }, { it.second }) + .mapValues { it.value.toSet() } + } + + fun getUnknownFeatures(unit: MLUnit<*>, featuresNames: Set): Set { + val knownFeatures = knownFeatures[unit] ?: return featuresNames + return featuresNames.filterNot { it in knownFeatures }.toSet() + } + } + + companion object { + private fun createKnownFeatures( + decisionFunction: DecisionFunctionWrapper, + modelUnits: Set>, + ): Map, FeatureFilter> { + fun createFeatureSelector(unit: MLUnit<*>) = object : FeatureFilter { + init { + val knownFeatures = decisionFunction.knownFeatures + knownFeatures.forEach { (unit, features) -> + val nonConsistentlyKnownFeatures = decisionFunction.getUnknownFeatures(unit, features) + require(nonConsistentlyKnownFeatures.isEmpty()) { + "These features are known and unknown at the same time: $nonConsistentlyKnownFeatures" + } + } + } + + override fun accept(featureDeclarations: Set>): Set> { + val availableFeaturesPerName = featureDeclarations.associateBy { it.name } + val availableFeaturesNames = featureDeclarations.map { it.name }.toSet() + val unknownFeaturesNames = decisionFunction.getUnknownFeatures(unit, availableFeaturesNames) + val knownAvailableFeaturesNames = availableFeaturesNames - unknownFeaturesNames + val knownAvailableFeatures = knownAvailableFeaturesNames.map { availableFeaturesPerName.getValue(it) }.toSet() + + return knownAvailableFeatures + } + + override fun accept(featureDeclaration: FeatureDeclaration<*>): Boolean { + val unknown = decisionFunction.getUnknownFeatures(unit, setOf(featureDeclaration.name)) + return unknown.isEmpty() + } + } + + return modelUnits.associateWith { createFeatureSelector(it) } + } + } +} diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/logs/TypingSpeedTracker.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/logs/TypingSpeedTracker.kt index 91771531c4a3..3f40516a7763 100644 --- a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/logs/TypingSpeedTracker.kt +++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/logs/TypingSpeedTracker.kt @@ -18,6 +18,8 @@ import kotlin.time.Duration import kotlin.time.Duration.Companion.seconds import kotlin.time.DurationUnit import com.intellij.platform.ml.feature.FeatureDeclaration as OldFeatureDeclaration +import com.jetbrains.ml.Feature as NewFeature +import com.jetbrains.ml.FeatureDeclaration as NewFeatureDeclaration @ApiStatus.Internal @Service @@ -49,6 +51,12 @@ class TypingSpeedTracker { } } + fun getTypingSpeedNewEventPairs(): Collection, NewFeature>> = DECAY_DURATIONS_NEW.mapNotNull { (decayDuration, eventFieldAndFeature) -> + typingSpeeds[decayDuration]?.let { + (eventFieldAndFeature.first with it) to (eventFieldAndFeature.second with it) + } + } + @TestOnly fun getTypingSpeed(decayDuration: Duration): Float? = typingSpeeds[decayDuration] @@ -88,9 +96,12 @@ class TypingSpeedTracker { ), OldFeatureDeclaration.float("typing_speed_${it}s").nullable()) } + private val DECAY_DURATIONS_NEW = listOf(1, 2, 5, 30) + .associate { it.seconds to Pair(EventFields.Float("typing_speed_${it}s"), NewFeatureDeclaration.float("typing_speed_${it}s").nullable()) } fun getInstance(): TypingSpeedTracker = service() fun getEventFields(): Array> = DECAY_DURATIONS.values.map { it.first }.toTypedArray() fun getFeatures(): Set> = DECAY_DURATIONS.values.map { it.second }.toSet() + fun getFeaturesNew(): List> = DECAY_DURATIONS_NEW.values.map { it.second } } } diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/TypingSpeedFeatureProvider.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/TypingSpeedFeatureProvider.kt new file mode 100644 index 000000000000..5eafa1ec2881 --- /dev/null +++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/TypingSpeedFeatureProvider.kt @@ -0,0 +1,26 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.intellij.codeInsight.inline.completion.ml + +import com.intellij.codeInsight.inline.completion.logs.TypingSpeedTracker +import com.jetbrains.ml.* + +internal class TypingSpeedFeatureProvider : FeatureProvider(MLUnitTyping) { + object Features { + val TIME_SINCE_LAST_TYPING = FeatureDeclaration.long("time_since_last_typing").nullable() + } + + override val featureComputationPolicy = FeatureComputationPolicy(true, true) + + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter): List = buildList { + with(units[MLUnitTyping]) { + val timeSinceLastTyping = getTimeSinceLastTyping() + if (timeSinceLastTyping != null) { + add(Features.TIME_SINCE_LAST_TYPING.with(timeSinceLastTyping)) + addAll(getTypingSpeedNewEventPairs().map { it.second }) + } + } + } + + + override val featureDeclarations = TypingSpeedTracker.getFeaturesNew() + Features.TIME_SINCE_LAST_TYPING +} diff --git a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/tiers.kt b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/tiers.kt index 6fbb4c11c85a..99465b1634c0 100644 --- a/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/tiers.kt +++ b/platform/platform-impl/src/com/intellij/codeInsight/inline/completion/ml/tiers.kt @@ -6,7 +6,7 @@ import com.intellij.lang.Language import com.intellij.openapi.editor.Editor import com.intellij.platform.ml.Tier import com.intellij.psi.PsiFile -import com.jetbrains.ml.features.api.MLUnit +import com.jetbrains.ml.MLUnit import org.jetbrains.annotations.ApiStatus @ApiStatus.Internal diff --git a/platform/platform-resources/src/META-INF/PlatformExtensionPoints.xml b/platform/platform-resources/src/META-INF/PlatformExtensionPoints.xml index ac2a8818765f..b6f1d1311abf 100644 --- a/platform/platform-resources/src/META-INF/PlatformExtensionPoints.xml +++ b/platform/platform-resources/src/META-INF/PlatformExtensionPoints.xml @@ -547,6 +547,9 @@ + diff --git a/platform/platform-resources/src/META-INF/PlatformExtensions.xml b/platform/platform-resources/src/META-INF/PlatformExtensions.xml index 67ee8ac03b33..4e2943147a86 100644 --- a/platform/platform-resources/src/META-INF/PlatformExtensions.xml +++ b/platform/platform-resources/src/META-INF/PlatformExtensions.xml @@ -1697,6 +1697,7 @@ + diff --git a/python/intellij.python.community.plugin.main.iml b/python/intellij.python.community.plugin.main.iml index 4e6198b803da..dae38e0f290d 100644 --- a/python/intellij.python.community.plugin.main.iml +++ b/python/intellij.python.community.plugin.main.iml @@ -11,6 +11,5 @@ - \ No newline at end of file diff --git a/python/intellij.python.community.plugin.modules.iml b/python/intellij.python.community.plugin.modules.iml index 433609c256d4..f67de7614cca 100644 --- a/python/intellij.python.community.plugin.modules.iml +++ b/python/intellij.python.community.plugin.modules.iml @@ -7,7 +7,6 @@ - diff --git a/python/intellij.python.ml.features/intellij.python.ml.features.iml b/python/intellij.python.ml.features/intellij.python.ml.features.iml deleted file mode 100644 index 9a7a08d0bac9..000000000000 --- a/python/intellij.python.ml.features/intellij.python.ml.features.iml +++ /dev/null @@ -1,23 +0,0 @@ - - - - - - - - - - - - - - - - - - - - - - - \ No newline at end of file diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/PyMLImportChooser.kt b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/PyMLImportChooser.kt deleted file mode 100644 index b569f80271ae..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/PyMLImportChooser.kt +++ /dev/null @@ -1,47 +0,0 @@ -package com.intellij.python.ml.features.imports - -import com.intellij.ide.DataManager -import com.intellij.openapi.actionSystem.DataContext -import com.intellij.openapi.application.ApplicationManager -import com.intellij.openapi.components.Service -import com.intellij.openapi.ui.popup.JBPopupFactory -import com.intellij.openapi.util.Computable -import com.intellij.util.Consumer -import com.jetbrains.python.codeInsight.imports.ImportCandidateHolder -import com.jetbrains.python.codeInsight.imports.PyImportChooser -import org.jetbrains.concurrency.AsyncPromise -import org.jetbrains.concurrency.Promise -import org.jetbrains.concurrency.resolvedPromise - -class PyMLImportChooser : PyImportChooser() { - override fun selectImport( - sources: MutableList, - useQualifiedImport: Boolean - ): Promise { - if (ApplicationManager.getApplication().isUnitTestMode()) { - return resolvedPromise(sources.get(0)) - } - val result = AsyncPromise() - - launchMLRanking(sources) { mlRanking: RateableRankingResult? -> - // GUI part - DataManager.getInstance().getDataContextFromFocus().doWhenDone(Consumer { dataContext: DataContext? -> - val popup = JBPopupFactory.getInstance() - .createPopupChooserBuilder(mlRanking!!.order) - .setItemChosenCallback(Consumer { item: ImportCandidateHolder? -> - result.setResult(item) - mlRanking.submitSelectedItem(item!!) - }) - .setCancelCallback(Computable { - mlRanking.submitPopUpClosed() - true - }) - processPopup(popup, useQualifiedImport) - popup.createPopup() - .showInBestPositionFor(dataContext!!) - }) - Unit - } - return result - } -} diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/README.md b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/README.md deleted file mode 100644 index 48ddd22995b2..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/README.md +++ /dev/null @@ -1,26 +0,0 @@ -# Ranking in fixing missing import quick fix - -## Registry - -* `quickfix.ranking.ml` - Status of the system - * `IN_EXPERIMENT` - ML is enabled depending on machine id. - * `ENABLED` - ML is enabled - * `DISABLED` - ML is disabled. ML features aren't computed -* `quickfix.ranking.ml.model.path` - Path to a local model - -## ML Model - -ML model loaded from resources of library `jetbrains.ml.models.python.imports.ranking.model` -(the dependency is located in `community/python/python-psi-impl/intellij.python.psi.impl.iml`). - -When you want, you can set `quickfix.ranking.ml.model.path` to an absolute path on your machine -and try out a freshly trained ML model. - -Note, that right now a regression model is used. -You may change that by tweaking `com.intellij.python.ml.features.imports.MLModelService`. - -## Contact - -Initial implementation: Gleb Marin -ML Model, A/B Experiment Pipelines: Nikita Ermolenko -PyCharm Support: Andrey Matveev diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlAnalysis.kt b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlAnalysis.kt deleted file mode 100644 index 42cd29997241..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlAnalysis.kt +++ /dev/null @@ -1,24 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.ml.features.imports - -import com.jetbrains.ml.features.api.logs.BooleanEventField -import com.jetbrains.ml.features.api.logs.EnumEventField -import com.jetbrains.ml.features.api.logs.IntEventField -import com.jetbrains.ml.features.api.logs.LongEventField - - -internal object ContextAnalysis { - val CANCELLED = BooleanEventField("selection_cancelled") { "No item has been selected" } - val SELECTED_POSITION = IntEventField("selected_position") { "The position of the selected import statement" } - val MODEL_UNAVAILABLE = BooleanEventField("model_unavailable") { "ML model was unavailable" } - val SELECTED_POSITION_INITIAL = IntEventField("selected_position_initial") { "The position of the selected import statement, if the final ML ranking would not have happened" } - val TIME_MS_TO_DISPLAY = LongEventField("time_ms_before_displayed") { "Duration from the quickfix start until when the imports were displayed" } - val TIME_MS_BEFORE_CLOSED = LongEventField("time_ms_before_closed") { "Duration from the quickfix start until the pop-up was closed" } - val ML_LOGGING_STATE = EnumEventField.of("ml_logging_state", { "State of the ML session logging" }) - val ML_ENABLED = BooleanEventField("ml_enabled") { "Machine Learning ranking is enabled" } - val ML_STATUS_CORRESPONDS_TO_BUCKET = BooleanEventField("ml_status_corresponds_to_bucket") { "Field 'ml_enabled' corresponds to the ML bucket % 2" } -} - -internal object CandidateAnalysis { - val MLFieldCorrectElement = BooleanEventField("is_correct") { "The candidate was chosen" } -} diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlImplementation.kt b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlImplementation.kt deleted file mode 100644 index 795b9d35d7bc..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlImplementation.kt +++ /dev/null @@ -1,202 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.ml.features.imports - -import com.intellij.openapi.application.writeAction -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.diagnostic.debug -import com.intellij.openapi.diagnostic.logger -import com.intellij.util.application -import com.jetbrains.ml.building.blocks.model.MLModel -import com.jetbrains.ml.building.blocks.task.MLFeaturesTree -import com.jetbrains.ml.features.api.logs.EventPair -import com.jetbrains.python.codeInsight.imports.ImportCandidateHolder -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportRankingContext -import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate -import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportRankingContext -import kotlinx.coroutines.CoroutineScope -import kotlinx.coroutines.Dispatchers -import kotlinx.coroutines.delay -import kotlinx.coroutines.launch -import java.util.concurrent.atomic.AtomicBoolean -import kotlin.time.Duration.Companion.seconds - - -@Service(Service.Level.APP) -internal class MLApiComputations( - val coroutineScope: CoroutineScope, -) - -internal sealed class FinalCandidatesRanker( - protected val contextFeatures: MLFeaturesTree, - protected val contextLogs: MutableList>, - protected val timestampStarted: Long, -) { - - abstract fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) - - abstract val mlEnabled: Boolean -} - -private class ExperimentalMLRanker( - mlContextFeatures: MLFeaturesTree, - mlContextLogs: MutableList>, - timestampStarted: Long, - private val mlModel: MLModel, -) : FinalCandidatesRanker(mlContextFeatures, mlContextLogs, timestampStarted) { - - override val mlEnabled = true - - override fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) { - service().coroutineScope.launch(Dispatchers.Default) { - contextFeatures.computeFeaturesFromProviders( - MLUnitImportRankingContext with ImportRankingContext(initialCandidatesOrder), - ) - - val importCandidatesFeatures = initialCandidatesOrder.associateWith { candidate -> - val candidateFeatures = contextFeatures.addChild() - candidateFeatures.computeFeaturesFromProviders( - MLUnitImportCandidate with ImportCandidateContext(initialCandidatesOrder, candidate) - ) - candidateFeatures - } - - val predictionsByCandidateTree = contextFeatures - .predictForChildren(mlModel, emptyMap()) - .toMap() - - val scoreByCandidate = initialCandidatesOrder.associateWith { candidate -> - val candidateTree = importCandidatesFeatures.getValue(candidate) - val score = predictionsByCandidateTree.getValue(candidateTree) - score - } - - val relevanceCandidateOrder = scoreByCandidate.toList() - .sortedByDescending { it.second } - .map { it.first } - - logger().debug { - "Sorted import candidates:\n" + - relevanceCandidateOrder.mapIndexed { index, candidate -> - val score = scoreByCandidate.getValue(candidate) - val initialIndex = initialCandidatesOrder.indexOf(candidate) - "\t#${index + 1}: '${candidate.presentableText}' (score: ${score}, initial index: $initialIndex)" - }.joinToString("\n") - } - - writeAction { - displayResult(RateableRankingResult(contextFeatures, contextLogs, importCandidatesFeatures, initialCandidatesOrder, relevanceCandidateOrder, timestampStarted)) - } - } - } -} - -private class InitialOrderKeepingRanker( - contextFeatures: MLFeaturesTree, - contextLogs: MutableList>, - timestampStarted: Long, -) : FinalCandidatesRanker(contextFeatures, contextLogs, timestampStarted) { - override val mlEnabled = false - - override fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) { - application.runWriteAction { - val importCandidatesFeatures = initialCandidatesOrder.associateWith { candidate -> - contextFeatures.addChild() - } - displayResult(RateableRankingResult(contextFeatures, contextLogs, importCandidatesFeatures, initialCandidatesOrder, initialCandidatesOrder, timestampStarted)) - } - } -} - - -internal fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) { - // In EDT - - val timestampStarted = System.currentTimeMillis() - return MLTaskPyCharmImportStatementsRanking.lifetime { - val mlContextFeatures = newMLFeaturesTree() - val mlContextLogs = mutableListOf>() - val rankingStatus = service().status - var modelWasReady = true - val ranker = if (rankingStatus.mlEnabled) { - val mlModel = service().getModel(MLTaskPyCharmImportStatementsRanking) - if (mlModel != null) { - ExperimentalMLRanker(mlContextFeatures, mlContextLogs, timestampStarted, mlModel) - } else { - modelWasReady = false - mlContextLogs.add(ContextAnalysis.MODEL_UNAVAILABLE with true) - InitialOrderKeepingRanker(mlContextFeatures, mlContextLogs, timestampStarted) - } - } - else { - InitialOrderKeepingRanker(mlContextFeatures, mlContextLogs, timestampStarted) - } - - mlContextLogs.add(ContextAnalysis.ML_STATUS_CORRESPONDS_TO_BUCKET with (rankingStatus.mlStatusCorrespondsToBucket && modelWasReady)) - mlContextLogs.add(ContextAnalysis.ML_ENABLED with (rankingStatus.mlEnabled && modelWasReady)) - - ranker.launchMLRanking(initialCandidatesOrder) { result -> - val timestampDisplayed = System.currentTimeMillis() - displayResult(result) - mlContextLogs.add(ContextAnalysis.TIME_MS_TO_DISPLAY with (timestampDisplayed - timestampStarted)) - } - } -} - -internal class RateableRankingResult( - private val mlContextFeatures: MLFeaturesTree, - private val contextAnalysis: MutableList>, - private val importCandidates: Map, - val orderInitial: List, - val order: List, - private val timestampStarted: Long, -) { - private var submitted = AtomicBoolean(false) - private val analysis: MutableMap>> = mutableMapOf( - mlContextFeatures to contextAnalysis - ) - - fun submitPopUpClosed() { - val timestampClosed = System.currentTimeMillis() - contextAnalysis.add(ContextAnalysis.TIME_MS_BEFORE_CLOSED with timestampClosed - timestampStarted) - service().coroutineScope.launch(Dispatchers.IO) { - delay(1.seconds) - synchronized(this@RateableRankingResult) { - if (submitted.get()) return@launch - submitted.set(true) - contextAnalysis.add(ContextAnalysis.CANCELLED with true) - } - submitLogs() - } - } - - fun submitSelectedItem(selected: ImportCandidateHolder) = synchronized(this) { - // In EDT - require(submitted.compareAndSet(false, true)) { "Some feedback to the ml ranker was already submitted" } - analysis[requireNotNull(importCandidates[selected])] = mutableListOf(CandidateAnalysis.MLFieldCorrectElement with true) - val selectedIndex = order.indexOf(selected) - val selectedIndexOriginal = orderInitial.indexOf(selected) - require(selectedIndex >= 0 && selectedIndexOriginal >= 0) - contextAnalysis.add(ContextAnalysis.SELECTED_POSITION with selectedIndex) - contextAnalysis.add(ContextAnalysis.SELECTED_POSITION_INITIAL with selectedIndexOriginal) - service().coroutineScope.launch(Dispatchers.IO) { - submitLogs() - } - } - - private suspend fun submitLogs() { - val logger = PyCharmImportsRankingLogs.mlLogger - val loggingOption = getLoggingOption(mlContextFeatures) - contextAnalysis.add(ContextAnalysis.ML_LOGGING_STATE with loggingOption) - when (loggingOption) { - LoggingOption.FULL -> logger.log(mlContextFeatures, analysis) - LoggingOption.NO_TREE -> { - val treeStem = logger.newMLLogsTree() - treeStem.addAnalysis(contextAnalysis) - logger.log(mlContextFeatures.sessionId, treeStem) - } - LoggingOption.SKIP -> Unit - } - } -} diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlLogs.kt b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlLogs.kt deleted file mode 100644 index f1d003eeb000..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlLogs.kt +++ /dev/null @@ -1,44 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.ml.features.imports - -import com.intellij.internal.statistic.eventLog.EventLogGroup -import com.intellij.internal.statistic.service.fus.collectors.CounterUsagesCollector -import com.intellij.platform.ml.impl.logs.MLEventLoggerProvider.Companion.ML_RECORDER_ID -import com.intellij.platform.ml.impl.logs.fus.IntelliJFusEventRegister -import com.intellij.util.application -import com.jetbrains.ml.building.blocks.logs.FusLoggerFactories -import com.jetbrains.ml.building.blocks.logs.extractEventFields -import com.jetbrains.ml.building.blocks.task.MLFeaturesTree - - -internal object PyCharmImportsRankingLogs : CounterUsagesCollector() { - private val GROUP = EventLogGroup("pycharm.quickfix.imports", 4, ML_RECORDER_ID) - val mlLogger = FusLoggerFactories - .withOneEvent( - fusEventName = "pycharm_import_statements_ranking", - fusEventRegister = IntelliJFusEventRegister(GROUP), - treeFeatures = MLTaskPyCharmImportStatementsRanking.treeFeatures, - treeAnalysisFields = listOf( - extractEventFields(ContextAnalysis), - extractEventFields(CandidateAnalysis) - ) - ) - - override fun getGroup() = GROUP -} - -internal enum class LoggingOption { - FULL, - NO_TREE, - SKIP -} - -internal const val RELEASE_FULL_LOGGING_PERCENT = 10 - -internal fun getLoggingOption(tree: MLFeaturesTree): LoggingOption { - if (application.isEAP || application.isUnitTestMode) { - return LoggingOption.FULL - } - if ((tree.sessionId * 37) % 100 > RELEASE_FULL_LOGGING_PERCENT) return LoggingOption.NO_TREE - return LoggingOption.FULL -} diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlModel.kt b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlModel.kt deleted file mode 100644 index 72de18fbaee8..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlModel.kt +++ /dev/null @@ -1,93 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.ml.features.imports - -import com.intellij.codeInsight.inline.completion.ml.MLUnitTyping -import com.intellij.openapi.components.Service -import com.intellij.openapi.components.service -import com.intellij.openapi.project.Project -import com.intellij.openapi.startup.ProjectActivity -import com.intellij.platform.ml.impl.MLModelSelector -import com.intellij.python.ml.features.imports.MissingTypingFeaturesLoader.Companion.fillingTypingFeatures -import com.jetbrains.ml.building.blocks.model.MLModel -import com.jetbrains.ml.building.blocks.model.MLModelLoader -import com.jetbrains.ml.building.blocks.model.catboost.DistributionFormat -import com.jetbrains.ml.building.blocks.model.catboost.ModelDistributionReaders -import com.jetbrains.ml.building.blocks.model.classical.MLModelLoaders -import com.jetbrains.ml.features.api.MLUnit -import com.jetbrains.ml.features.api.feature.Feature -import com.jetbrains.ml.features.api.feature.FeatureDeclaration -import java.nio.file.Path - - -@Service -class MLModelService : MLModelSelector, Double>( - defaultLoader = MLModelLoaders.Regressor( - reader = ModelDistributionReaders.FromJavaResources(MLModelService::class.java, Path.of("python-imports-ranking-model")), - format = DistributionFormat(), - ).fillingTypingFeatures(), - registryModelLoaderFactory = { path -> - MLModelLoaders.Regressor( - reader = ModelDistributionReaders.FromDirectory(path), - format = DistributionFormat(), - ).fillingTypingFeatures() - }, - modelPathRegistryKey = "quickfix.ranking.ml.model.path", - coroutineScope = service().coroutineScope, -) - -private class MissingTypingFeaturesLoader(private val baseLoader: MLModelLoader, Double>) : MLModelLoader, Double> by baseLoader { - override suspend fun loadModel(features: Map, List>>, parameters: Map?): MLModel { - val model = baseLoader.loadModel(features + mapOf( - MLUnitTyping to listOf( - TypingFeatures.SINCE_LAST_TYPING, - TypingFeatures.TYPING_SPEED_1S, - TypingFeatures.TYPING_SPEED_2S, - TypingFeatures.TYPING_SPEED_30S, - TypingFeatures.TYPING_SPEED_5S - ) - ), parameters) - - return object : MLModel by model { - override fun predict(features: Map, List>, parameters: Map): Double { - return model.predict(features + mapOf(MLUnitTyping to fillTypingFeatures()), parameters) - } - - override fun predictBatch(contextFeatures: Map, List>, itemFeatures: List, List>>, parameters: Map): List { - return model.predictBatch(contextFeatures + mapOf(MLUnitTyping to fillTypingFeatures()), itemFeatures, parameters) - } - } - } - - companion object { - fun MLModelLoader, Double>.fillingTypingFeatures(): MLModelLoader, Double> { - return MissingTypingFeaturesLoader(this) - } - } -} - -/** - * TODO: Should be removed, because they don't exist anymore - */ -internal object TypingFeatures { - val SINCE_LAST_TYPING = FeatureDeclaration.int("time_since_last_typing") - val TYPING_SPEED_1S = FeatureDeclaration.double("typing_speed_1s") - val TYPING_SPEED_2S = FeatureDeclaration.double("typing_speed_2s") - val TYPING_SPEED_30S = FeatureDeclaration.double("typing_speed_30s") - val TYPING_SPEED_5S = FeatureDeclaration.double("typing_speed_5s") -} - -private class QuickfixRankingModelLoading : ProjectActivity { - override suspend fun execute(project: Project) { - service().prepareModel(MLTaskPyCharmImportStatementsRanking) - } -} - -private fun fillTypingFeatures(): List { - return listOf( - TypingFeatures.SINCE_LAST_TYPING with 1, - TypingFeatures.TYPING_SPEED_2S with 0.0, - TypingFeatures.TYPING_SPEED_30S with 0.0, - TypingFeatures.TYPING_SPEED_5S with 0.0, - TypingFeatures.TYPING_SPEED_1S with 0.0, - ) -} diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlTask.kt b/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlTask.kt deleted file mode 100644 index 127fb298b467..000000000000 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlTask.kt +++ /dev/null @@ -1,29 +0,0 @@ -// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.ml.features.imports - -import com.jetbrains.ml.building.blocks.task.MLTask -import com.jetbrains.python.codeInsight.imports.mlapi.features.* - - -val MLTaskPyCharmImportStatementsRanking = MLTask( - id = "pycharm_import_statements_ranking", - treeFeaturesPerLevel = listOf( - { - featuresFromProviders( - CandidatesListFeatures, - BaseProjectFeatures - ) - }, - { - featuresFromProviders( - ImportsFeatures, - RelevanceEvaluationFeatures, - ImportCandidateRelativeFeatures, - NeighborFilesImportsFeatures, - PrimitiveImportFeatures, - PsiStructureFeatures, - OpenFilesImportsFeatures, - ) - } - ) -) diff --git a/python/intellij.python.ml.features/src/intellij.python.ml.features.xml b/python/intellij.python.ml.features/src/intellij.python.ml.features.xml deleted file mode 100644 index 9b29e3da5ae6..000000000000 --- a/python/intellij.python.ml.features/src/intellij.python.ml.features.xml +++ /dev/null @@ -1,14 +0,0 @@ - - - - - - - - diff --git a/python/pluginCore/plugin-content.yaml b/python/pluginCore/plugin-content.yaml index 2d6e11342590..2e7919ffef7b 100644 --- a/python/pluginCore/plugin-content.yaml +++ b/python/pluginCore/plugin-content.yaml @@ -4,10 +4,6 @@ files: - name: $MAVEN_REPOSITORY$/org/apache/ws/xmlrpc/xmlrpc/2/xmlrpc-2.jar reason: withProjectLibrary - - name: jetbrains.ml.models.python.imports.ranking.model - files: - - name: $MAVEN_REPOSITORY$/com/jetbrains/ml/models/python-imports-ranking-model/lilac-coua/python-imports-ranking-model-lilac-coua.jar - reason: <- intellij.python.ml.features - name: jsr305 files: - name: $MAVEN_REPOSITORY$/com/google/code/findbugs/jsr305/3/jsr305-3.jar @@ -46,7 +42,6 @@ - name: intellij.python.grazie - name: intellij.python.langInjection - name: intellij.python.markdown - - name: intellij.python.ml.features - name: intellij.python.terminal - name: lib/python-common.jar modules: diff --git a/python/pluginCore/resources/META-INF/plugin.xml b/python/pluginCore/resources/META-INF/plugin.xml index c26acb7b06a6..7704e173f3bf 100644 --- a/python/pluginCore/resources/META-INF/plugin.xml +++ b/python/pluginCore/resources/META-INF/plugin.xml @@ -52,7 +52,6 @@ The Python plug-in provides smart editing for Python scripts. The feature set of - @@ -522,6 +521,7 @@ The Python plug-in provides smart editing for Python scripts. The feature set of + + + + + + + + + + + + diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/BaseProjectFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/BaseProjectFeatures.kt index f0e19b892b18..b0059e3dd9bc 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/BaseProjectFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/BaseProjectFeatures.kt @@ -10,10 +10,9 @@ import com.intellij.psi.PsiManager import com.intellij.psi.search.FileTypeIndex import com.intellij.psi.search.GlobalSearchScope import com.intellij.psi.util.PsiTreeUtil -import com.jetbrains.ml.features.api.feature.* +import com.jetbrains.ml.* import com.jetbrains.python.PythonFileType -import com.jetbrains.python.codeInsight.imports.mlapi.ImportRankingContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportRankingContextFeatures +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidatesList import com.jetbrains.python.psi.* private val interestingClasses = arrayOf( @@ -32,7 +31,7 @@ enum class FileExtensionType { PXI // .pxi files } -object BaseProjectFeatures : ImportRankingContextFeatures() { +class BaseProjectFeatures : FeatureProvider(MLUnitImportCandidatesList) { object Features { val NUM_PYTHON_FILES_IN_PROJECT = FeatureDeclaration.int("num_python_files_in_project") { "The estimated amount of files in the project (by a power of 2)" @@ -42,10 +41,10 @@ object BaseProjectFeatures : ImportRankingContextFeatures() { } override val featureComputationPolicy = FeatureComputationPolicy(false, true) - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) - override suspend fun computeFeatures(instance: ImportRankingContext, filter: FeatureFilter): List? = buildList { - val candidates = instance.candidates + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter): List = buildList { + val candidates = units[MLUnitImportCandidatesList] if (candidates.isEmpty()) return@buildList val project = candidates[0].importable?.project ?: return@buildList diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/CandidatesListFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/CandidatesListFeatures.kt index 420f762a9d16..d1c4e3f3df40 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/CandidatesListFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/CandidatesListFeatures.kt @@ -1,14 +1,10 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.codeInsight.imports.mlapi.features -import com.jetbrains.ml.features.api.feature.Feature -import com.jetbrains.ml.features.api.feature.FeatureDeclaration -import com.jetbrains.ml.features.api.feature.FeatureFilter -import com.jetbrains.ml.features.api.feature.extractFeatureDeclarations -import com.jetbrains.python.codeInsight.imports.mlapi.ImportRankingContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportRankingContextFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidatesList -object CandidatesListFeatures : ImportRankingContextFeatures() { +class CandidatesListFeatures : FeatureProvider(MLUnitImportCandidatesList) { object Features { val LENGTH = FeatureDeclaration.int("n_candidates") { "The amount of import candidates" @@ -18,10 +14,11 @@ object CandidatesListFeatures : ImportRankingContextFeatures() { }.nullable() } - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) - override suspend fun computeFeatures(instance: ImportRankingContext, filter: FeatureFilter): List = buildList { - add(Features.LENGTH with instance.candidates.size) - add(Features.HIGHEST_OLD_RELEVANCE with instance.candidates.maxOfOrNull { it.relevance }) + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) = buildList { + val candidates = units[MLUnitImportCandidatesList] + add(Features.LENGTH with candidates.size) + add(Features.HIGHEST_OLD_RELEVANCE with candidates.maxOfOrNull { it.relevance }) } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportCandidateRelativeFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportCandidateRelativeFeatures.kt index ce37714a9a14..a7ecf903b96c 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportCandidateRelativeFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportCandidateRelativeFeatures.kt @@ -1,23 +1,27 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.codeInsight.imports.mlapi.features -import com.jetbrains.ml.features.api.feature.Feature -import com.jetbrains.ml.features.api.feature.FeatureDeclaration -import com.jetbrains.ml.features.api.feature.FeatureFilter -import com.jetbrains.ml.features.api.feature.extractFeatureDeclarations -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidatesList -object ImportCandidateRelativeFeatures : ImportCandidateFeatures() { +class ImportCandidateRelativeFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val RELATIVE_POSITION = FeatureDeclaration.int("original_position") { "The import's position without ML ranking" } } - override val featureDeclarations = extractFeatureDeclarations(Features) + override val unitSight = setOf( + MLUnitImportCandidatesList + ) - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List = buildList { - add(Features.RELATIVE_POSITION with instance.candidates.indexOf(instance.candidate)) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) + + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) = buildList { + val candidate = units[MLUnitImportCandidate] + val allCandidates = units[MLUnitImportCandidatesList] + + add(Features.RELATIVE_POSITION with allCandidates.indexOf(candidate)) } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportsFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportsFeatures.kt index d4437a874867..4ca120786284 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportsFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/ImportsFeatures.kt @@ -5,16 +5,15 @@ import com.intellij.openapi.application.readAction import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.psi.PsiFile import com.intellij.psi.PsiManager -import com.intellij.psi.util.QualifiedName -import com.jetbrains.ml.features.api.feature.* -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate import com.jetbrains.python.psi.PyFile +import one.util.streamex.StreamEx +import com.intellij.psi.util.QualifiedName import com.jetbrains.python.psi.PyFromImportStatement import com.jetbrains.python.psi.PyImportElement -import one.util.streamex.StreamEx -object ImportsFeatures : ImportCandidateFeatures() { +class ImportsFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val EXISTING_IMPORT_FROM_PREFIX = FeatureDeclaration.int("existing_import_from_prefix") { "A maximal prefix for which there exists some import from" }.nullable() val EXISTING_IMPORT_PREFIX = FeatureDeclaration.int("existing_import_prefix") { "A maximal prefix for which there exists some import" }.nullable() @@ -24,11 +23,12 @@ object ImportsFeatures : ImportCandidateFeatures() { override val featureComputationPolicy = FeatureComputationPolicy(false, true) - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) + + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) = buildList { + val importCandidate = units[MLUnitImportCandidate] + val project = importCandidate.importable?.project ?: return@buildList - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List = buildList { - val importCandidate = instance.candidate - val project = readAction { importCandidate.importable?.project } ?: return@buildList val editor = readAction { FileEditorManager.getInstance(project).selectedEditor } ?: return@buildList val psiFile = readAction { PsiManager.getInstance(project).findFile(editor.file) } ?: return@buildList diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/NeighborFilesImportsFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/NeighborFilesImportsFeatures.kt index 5cb0163c5ed0..6b1f06f9169f 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/NeighborFilesImportsFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/NeighborFilesImportsFeatures.kt @@ -5,16 +5,15 @@ import com.intellij.openapi.application.readAction import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.openapi.vfs.VirtualFile import com.intellij.psi.PsiManager -import com.jetbrains.ml.features.api.feature.* -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate import com.jetbrains.python.psi.PyFile import kotlinx.coroutines.async import kotlinx.coroutines.coroutineScope private const val FILES_TO_WATCH = 8 -object NeighborFilesImportsFeatures : ImportCandidateFeatures() { +class NeighborFilesImportsFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val NEIGHBOR_FILES_EXISTING_IMPORT_FROM_PREFIX = FeatureDeclaration.int("neighbor_files_existing_import_from_prefix") { "A maximal prefix for which there exists some import from" }.nullable() val NEIGHBOR_FILES_EXISTING_IMPORT_PREFIX = FeatureDeclaration.int("neighbor_files_existing_import_prefix") { "A maximal prefix for which there exists some import" }.nullable() @@ -24,10 +23,10 @@ object NeighborFilesImportsFeatures : ImportCandidateFeatures() { override val featureComputationPolicy = FeatureComputationPolicy(false, true) - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List = coroutineScope { - val importCandidate = instance.candidate + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter): List = coroutineScope { + val importCandidate = units[MLUnitImportCandidate] if (importCandidate.path == null) { return@coroutineScope buildList { add(Features.NEIGHBOR_FILES_EXISTING_IMPORT_FROM_PREFIX with 0) diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/OpenFilesImportsFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/OpenFilesImportsFeatures.kt index 3bc44b4b5812..6f53b6c7ffaf 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/OpenFilesImportsFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/OpenFilesImportsFeatures.kt @@ -4,16 +4,15 @@ package com.jetbrains.python.codeInsight.imports.mlapi.features import com.intellij.openapi.application.readAction import com.intellij.openapi.fileEditor.FileEditorManager import com.intellij.psi.PsiManager -import com.jetbrains.ml.features.api.feature.* -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate import com.jetbrains.python.psi.PyFile -import kotlinx.coroutines.async -import kotlinx.coroutines.coroutineScope +import kotlinx.coroutines.* +import kotlin.collections.copyOfRange private const val FILES_TO_WATCH = 8 -object OpenFilesImportsFeatures : ImportCandidateFeatures() { +class OpenFilesImportsFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val OPEN_FILES_EXISTING_IMPORT_FROM_PREFIX = FeatureDeclaration.int("open_files_existing_import_from_prefix") { "A maximal prefix for which there exists some import from" }.nullable() val OPEN_FILES_EXISTING_IMPORT_PREFIX = FeatureDeclaration.int("open_files_existing_import_prefix") { "A maximal prefix for which there exists some import" }.nullable() @@ -23,10 +22,10 @@ object OpenFilesImportsFeatures : ImportCandidateFeatures() { override val featureComputationPolicy = FeatureComputationPolicy(false, true) - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List = coroutineScope { - val importCandidate = instance.candidate + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) : List = coroutineScope { + val importCandidate = units[MLUnitImportCandidate] if (importCandidate.path == null) { return@coroutineScope buildList { add(Features.OPEN_FILES_EXISTING_IMPORT_FROM_PREFIX with 0) diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PrimitiveImportFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PrimitiveImportFeatures.kt index 9c54a5d4489c..211a65530b07 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PrimitiveImportFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PrimitiveImportFeatures.kt @@ -1,15 +1,11 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.codeInsight.imports.mlapi.features -import com.jetbrains.ml.features.api.feature.Feature -import com.jetbrains.ml.features.api.feature.FeatureDeclaration -import com.jetbrains.ml.features.api.feature.FeatureFilter -import com.jetbrains.ml.features.api.feature.extractFeatureDeclarations -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate -object PrimitiveImportFeatures : ImportCandidateFeatures() { +class PrimitiveImportFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val RELEVANCE = FeatureDeclaration.int("old_relevance") { """ @@ -21,10 +17,13 @@ object PrimitiveImportFeatures : ImportCandidateFeatures() { }.nullable() } - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List? = buildList { - add(Features.RELEVANCE with instance.candidate.relevance) - add(Features.COMPONENT_COUNT with instance.candidate.path?.componentCount) + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) = buildList { + + val importCandidate = units[MLUnitImportCandidate] + + add(Features.RELEVANCE with importCandidate.relevance) + add(Features.COMPONENT_COUNT with importCandidate.path?.componentCount) } } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PsiStructureFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PsiStructureFeatures.kt index 0be01b83d3c6..1d62a44ee936 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PsiStructureFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/PsiStructureFeatures.kt @@ -3,28 +3,26 @@ package com.jetbrains.python.codeInsight.imports.mlapi.features import com.intellij.openapi.application.readAction import com.intellij.psi.PsiElement -import com.jetbrains.ml.features.api.feature.Feature -import com.jetbrains.ml.features.api.feature.FeatureDeclaration -import com.jetbrains.ml.features.api.feature.FeatureFilter -import com.jetbrains.ml.features.api.feature.extractFeatureDeclarations -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.ImportCandidateHolder +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate -object PsiStructureFeatures : ImportCandidateFeatures() { +class PsiStructureFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val PSI_CLASS = FeatureDeclaration.aClass("importable_class") { "PSI class of the imported element" }.nullable() val PSI_PARENT = (1..4).map { i -> FeatureDeclaration.aClass("psi_parent_$i") { "PSI parent #$i" }.nullable() } } - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) = buildList { + val importCandidate: ImportCandidateHolder = units[MLUnitImportCandidate] - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List = buildList { readAction { - add(Features.PSI_CLASS with (instance.candidate.importable?.javaClass)) + add(Features.PSI_CLASS with (importCandidate.importable?.javaClass)) Features.PSI_PARENT.withIndex().forEach { (i, featureDeclaration) -> - var parent: PsiElement? = instance.candidate.importable + var parent: PsiElement? = importCandidate.importable repeat(i) { parent = parent?.parent } diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/RelevanceEvaluationFeatures.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/RelevanceEvaluationFeatures.kt index 44085000291f..c2b999977bb8 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/RelevanceEvaluationFeatures.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/features/RelevanceEvaluationFeatures.kt @@ -6,9 +6,8 @@ import com.intellij.openapi.module.ModuleUtilCore import com.intellij.openapi.projectRoots.Sdk import com.intellij.openapi.vfs.VirtualFile import com.intellij.psi.util.QualifiedName -import com.jetbrains.ml.features.api.feature.* -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateContext -import com.jetbrains.python.codeInsight.imports.mlapi.ImportCandidateFeatures +import com.jetbrains.ml.* +import com.jetbrains.python.codeInsight.imports.mlapi.MLUnitImportCandidate import com.jetbrains.python.sdk.PythonSdkUtil enum class UnderscoresType { @@ -21,7 +20,6 @@ enum class UnderscoresType { DOUBLE_LEADING_TRAILING, // __something__ IRREGULAR // not from the above } - enum class ModuleSourceType { STD_LIB, EXTERNAL_LIB, @@ -29,7 +27,9 @@ enum class ModuleSourceType { } -object RelevanceEvaluationFeatures : ImportCandidateFeatures() { + + +class RelevanceEvaluationFeatures : FeatureProvider(MLUnitImportCandidate) { object Features { val UNDERSCORES_IN_PATH = FeatureDeclaration.int("underscores_in_path") { "number of prefix and suffix underscores in path" }.nullable() val MODULE_SOURCE_TYPE = FeatureDeclaration.enum("module_source_type") { "info about lib being std, local, or external" }.nullable() @@ -38,10 +38,10 @@ object RelevanceEvaluationFeatures : ImportCandidateFeatures() { override val featureComputationPolicy = FeatureComputationPolicy(false, true) - override val featureDeclarations = extractFeatureDeclarations(Features) + override val featureDeclarations = extractFieldsAsFeatureDeclarations(Features) - override suspend fun computeFeatures(instance: ImportCandidateContext, filter: FeatureFilter): List = buildList { - val importCandidate = instance.candidate + override suspend fun computeFeatures(units: MLUnitsMap, usefulFeaturesFilter: FeatureFilter) = buildList { + val importCandidate = units[MLUnitImportCandidate] add(Features.UNDERSCORES_IN_PATH with countBoundaryUnderscores(importCandidate.path)) readAction { Features.UNDERSCORES_TYPES_OF_PACKAGES.withIndex().forEach { (i, featureDeclaration) -> @@ -64,10 +64,7 @@ object RelevanceEvaluationFeatures : ImportCandidateFeatures() { ModuleUtilCore.findModuleForFile(vFile, baseElement.project) == null -> ModuleSourceType.EXTERNAL_LIB else -> ModuleSourceType.LOCAL_LIB } - } - ) - } else { - add(Features.MODULE_SOURCE_TYPE with null) + }) } } @@ -75,7 +72,6 @@ object RelevanceEvaluationFeatures : ImportCandidateFeatures() { if (qName == null) return 0 return qName.components.sumOf { it.takeWhile { it == '_' }.length + it.takeLastWhile { it == '_' }.length } } - private fun getUnderscoresType(s: String): UnderscoresType { val leadingUnderscores = s.takeWhile { it == '_' }.length val trailingUnderscores = s.reversed().takeWhile { it == '_' }.length diff --git a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/mlUnits.kt b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/mlUnits.kt index 88560cb51884..4d686b2d26cf 100644 --- a/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/mlUnits.kt +++ b/python/python-psi-impl/src/com/jetbrains/python/codeInsight/imports/mlapi/mlUnits.kt @@ -1,25 +1,9 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. package com.jetbrains.python.codeInsight.imports.mlapi -import com.jetbrains.ml.features.api.MLUnit -import com.jetbrains.ml.features.api.feature.FeatureProvider +import com.jetbrains.ml.MLUnit import com.jetbrains.python.codeInsight.imports.ImportCandidateHolder +val MLUnitImportCandidate = MLUnit("import_candidate") -data class ImportRankingContext( - val candidates: List -) - -data class ImportCandidateContext( - val candidates: List, - val candidate: ImportCandidateHolder, -) - -val MLUnitImportRankingContext = MLUnit("import_candidates_list") - -val MLUnitImportCandidate = MLUnit("import_candidate") - - -abstract class ImportRankingContextFeatures : FeatureProvider(MLUnitImportRankingContext) - -abstract class ImportCandidateFeatures : FeatureProvider(MLUnitImportCandidate) +val MLUnitImportCandidatesList = MLUnit>("import_candidates_list") diff --git a/python/src/com/jetbrains/python/codeInsight/imports/PyImportChooser.java b/python/src/com/jetbrains/python/codeInsight/imports/PyImportChooser.java index 61168a4d4a3e..7a55a18b53f2 100644 --- a/python/src/com/jetbrains/python/codeInsight/imports/PyImportChooser.java +++ b/python/src/com/jetbrains/python/codeInsight/imports/PyImportChooser.java @@ -7,13 +7,13 @@ import com.intellij.openapi.application.ApplicationManager; import com.intellij.openapi.editor.colors.EditorColorsManager; import com.intellij.openapi.editor.colors.EditorColorsScheme; import com.intellij.openapi.editor.colors.EditorFontType; -import com.intellij.openapi.ui.popup.IPopupChooserBuilder; import com.intellij.openapi.ui.popup.JBPopupFactory; import com.intellij.psi.PsiElement; import com.intellij.ui.SimpleColoredComponent; import com.intellij.ui.SimpleTextAttributes; import com.intellij.util.Consumer; import com.jetbrains.python.PyPsiBundle; +import kotlin.Unit; import org.jetbrains.concurrency.AsyncPromise; import org.jetbrains.concurrency.Promise; import org.jetbrains.concurrency.Promises; @@ -22,34 +22,39 @@ import javax.swing.*; import java.awt.*; import java.util.List; -public class PyImportChooser implements ImportChooser { +import static com.jetbrains.python.codeInsight.imports.mlapi.MlImplementationKt.launchMLRanking; + + +public final class PyImportChooser implements ImportChooser { @Override public Promise selectImport(List sources, boolean useQualifiedImport) { if (ApplicationManager.getApplication().isUnitTestMode()) { return Promises.resolvedPromise(sources.get(0)); } - AsyncPromise result = new AsyncPromise<>(); - // GUI part - DataManager.getInstance().getDataContextFromFocus().doWhenDone((Consumer)dataContext -> { - var popup = JBPopupFactory.getInstance() - .createPopupChooserBuilder(sources) - .setItemChosenCallback(item -> { - result.setResult(item); - }); - processPopup(popup, useQualifiedImport); - popup.createPopup() - .showInBestPositionFor(dataContext); + launchMLRanking(sources, (mlRanking) -> { + // GUI part + DataManager.getInstance().getDataContextFromFocus().doWhenDone((Consumer)dataContext -> JBPopupFactory.getInstance() + .createPopupChooserBuilder(mlRanking.getOrder()) + .setRenderer(new CellRenderer()) + .setTitle(useQualifiedImport ? PyPsiBundle.message("ACT.qualify.with.module") : PyPsiBundle.message("ACT.from.some.module.import")) + .setItemChosenCallback(item -> { + result.setResult(item); + mlRanking.submitSelectedItem(item); + }) + .setCancelCallback(() -> { + mlRanking.submitPopUpClosed(); + return true; + }) + .setNamerForFiltering(o -> o.getPresentableText()) + .createPopup() + .showInBestPositionFor(dataContext)); + return Unit.INSTANCE; }); - return result; - } - protected void processPopup(IPopupChooserBuilder popup, boolean useQualifiedImport) { - popup.setRenderer(new CellRenderer()) - .setTitle(useQualifiedImport ? PyPsiBundle.message("ACT.qualify.with.module") : PyPsiBundle.message("ACT.from.some.module.import")) - .setNamerForFiltering(o -> o.getPresentableText()); + return result; } // Stolen from FQNameCellRenderer diff --git a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlAB.kt b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlAB.kt similarity index 96% rename from python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlAB.kt rename to python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlAB.kt index aff7c18b8e10..4dfddac55351 100644 --- a/python/intellij.python.ml.features/src/com/intellij/python/ml/features/imports/mlAB.kt +++ b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlAB.kt @@ -1,5 +1,5 @@ // Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. -package com.intellij.python.ml.features.imports +package com.jetbrains.python.codeInsight.imports.mlapi import com.intellij.internal.statistic.eventLog.EventLogConfiguration import com.intellij.openapi.components.Service diff --git a/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlAnalysis.kt b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlAnalysis.kt new file mode 100644 index 000000000000..afa94199d656 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlAnalysis.kt @@ -0,0 +1,42 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.imports.mlapi + +import com.jetbrains.ml.analysis.MLTaskFeedback +import com.jetbrains.ml.analysis.MLTaskFeedbackField +import com.jetbrains.ml.analysis.MLUnitFeedbackField +import com.jetbrains.ml.logs.schema.BooleanEventField +import com.jetbrains.ml.logs.schema.EnumEventField +import com.jetbrains.ml.logs.schema.IntEventField +import com.jetbrains.ml.logs.schema.LongEventField + +// suspendable analysis -- written after the imports were ranked + +internal val MLFieldCorrectElement = MLUnitFeedbackField( + unit = MLUnitImportCandidate, + field = BooleanEventField("is_correct") { "The candidate was chosen" }, + throwOnTimeout = false +) + + +internal object MLFeedbackCorrectElementPosition : MLTaskFeedback() { + val CANCELLED = BooleanEventField("selection_cancelled") { "No item has been selected" } + val SELECTED_POSITION = IntEventField("selected_position") { "The position of the selected import statement" } + val INVALID_SELECTED_ITEM = BooleanEventField("invalid_selected_item") { "The selected item hasn't been found in the initial list" } + + override val feedbackDeclaration = listOf(CANCELLED, SELECTED_POSITION, INVALID_SELECTED_ITEM) +} + +internal val ML_FEEDBACK_TIME_MS_TO_DISPLAY = MLTaskFeedbackField( + field = LongEventField("time_ms_before_displayed") { "Duration from the quickfix start until when the imports were displayed" }, + throwOnTimeout = true, +) + +internal val ML_FEEDBACK_TIME_MS_BEFORE_CLOSED = MLTaskFeedbackField( + field = LongEventField("time_ms_before_closed") { "Duration from the quickfix start until the pop-up was closed" }, +) + +// runtime analysis -- written during the imports are ranked + +internal val ML_LOGGING_STATE: EnumEventField = EnumEventField.of("ml_logging_state", { "State of the ML session logging" }) +internal val ML_ENABLED = BooleanEventField("ml_enabled") { "Machine Learning ranking is enabled" } +internal val ML_STATUS_CORRESPONDS_TO_BUCKET = BooleanEventField("ml_status_corresponds_to_bucket") { "Field 'ml_enabled' corresponds to the ML bucket % 2" } diff --git a/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlImplementation.kt b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlImplementation.kt new file mode 100644 index 000000000000..ad52dd2cd917 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlImplementation.kt @@ -0,0 +1,136 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.imports.mlapi + +import com.intellij.codeInsight.inline.completion.logs.TypingSpeedTracker +import com.intellij.codeInsight.inline.completion.ml.MLUnitTyping +import com.intellij.openapi.application.writeAction +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.util.application +import com.jetbrains.ml.MockMLModel +import com.jetbrains.ml.platform.MLApiTaskExecutor +import com.jetbrains.ml.session.MLSession +import com.jetbrains.python.codeInsight.imports.ImportCandidateHolder +import kotlinx.coroutines.CoroutineScope +import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.delay +import kotlinx.coroutines.launch +import java.util.concurrent.atomic.AtomicBoolean +import kotlin.time.Duration.Companion.seconds + + +@Service(Service.Level.APP) +private class MLApiComputations( + val coroutineScope: CoroutineScope, +) + +internal sealed class FinalCandidatesRanker( + protected val mlSession: MLSession, + protected val timestampStarted: Long, +) { + + abstract fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) + + abstract val mlEnabled: Boolean +} + +private class ExperimentalMLRanker( + private val coroutineScope: CoroutineScope, + mlSession: MLSession, timestampStarted: Long, +) : FinalCandidatesRanker(mlSession, timestampStarted) { + + override val mlEnabled = true + + override fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) { + coroutineScope.launch(Dispatchers.Default) { + val mlTreeRoot = mlSession.buildRoot( + MLUnitImportCandidatesList with initialCandidatesOrder, + MLUnitTyping with TypingSpeedTracker.getInstance(), + ).computingFeatures { computeAllFeaturesNow() } + + for (candidate in initialCandidatesOrder) { + mlTreeRoot + .onlyLog(MLUnitImportCandidate with candidate) + .computingFeatures { computeAllFeaturesNow() } + } + + mlSession.finish() + + writeAction { + displayResult(RateableRankingResult(mlSession, initialCandidatesOrder, timestampStarted)) + } + } + } +} + +private class InitialOrderKeepingRanker( + mlSession: MLSession, timestampStarted: Long, +) : FinalCandidatesRanker(mlSession, timestampStarted) { + override val mlEnabled = false + + override fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) { + mlSession.finish() + + application.runWriteAction { + displayResult(RateableRankingResult(mlSession, initialCandidatesOrder, timestampStarted)) + } + } +} + + +internal fun launchMLRanking(initialCandidatesOrder: MutableList, displayResult: (RateableRankingResult) -> Unit) { + // In EDT + + val timestampStarted = System.currentTimeMillis() + val mlCoroutineScope = service().coroutineScope + val mlSession = service().task + .startMLSession(taskExecutor = MLApiTaskExecutor.configure(mlCoroutineScope)) + val rankingStatus = service().status + val ranker = if (rankingStatus.mlEnabled) + ExperimentalMLRanker(mlCoroutineScope, mlSession, timestampStarted) + else + InitialOrderKeepingRanker(mlSession, timestampStarted) + + mlSession.writeRuntimeSessionAnalysis(ML_LOGGING_STATE with getLoggingOption(mlSession)) + mlSession.writeRuntimeSessionAnalysis(ML_ENABLED with rankingStatus.mlEnabled) + mlSession.writeRuntimeSessionAnalysis(ML_STATUS_CORRESPONDS_TO_BUCKET with rankingStatus.mlStatusCorrespondsToBucket) + + ranker.launchMLRanking(initialCandidatesOrder) { result -> + val timestampDisplayed = System.currentTimeMillis() + displayResult(result) + ML_FEEDBACK_TIME_MS_TO_DISPLAY.feedback(mlSession, timestampDisplayed - timestampStarted) + } +} + +internal class RateableRankingResult( + private val mlSession: MLSession<*, *>, + val order: List, + private val timestampStarted: Long, +) { + private var submitted = AtomicBoolean(false) + + fun submitPopUpClosed() { + val timestampClosed = System.currentTimeMillis() + ML_FEEDBACK_TIME_MS_BEFORE_CLOSED.feedback(mlSession, timestampClosed - timestampStarted) + service().coroutineScope.launch { + delay(1.seconds) + synchronized(this@RateableRankingResult) { + if (submitted.get()) return@launch + submitted.set(true) + MLFeedbackCorrectElementPosition.feedbackEventPairs(mlSession, MLFeedbackCorrectElementPosition.CANCELLED with true) + MLFieldCorrectElement.cancelFeedback(mlSession) + } + } + } + + fun submitSelectedItem(selected: ImportCandidateHolder) = synchronized(this) { + // In EDT + require(submitted.compareAndSet(false, true)) { "Some feedback to the ml ranker was already submitted" } + MLFieldCorrectElement.feedback(mlSession, selected, true) + val selectedIndex = order.indexOf(selected) + if (selectedIndex == -1) + MLFeedbackCorrectElementPosition.feedbackEventPairs(mlSession, MLFeedbackCorrectElementPosition.INVALID_SELECTED_ITEM with true) + else + MLFeedbackCorrectElementPosition.feedbackEventPairs(mlSession, MLFeedbackCorrectElementPosition.SELECTED_POSITION with selectedIndex) + } +} diff --git a/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlLogs.kt b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlLogs.kt new file mode 100644 index 000000000000..4d1b54300b52 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlLogs.kt @@ -0,0 +1,80 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.imports.mlapi + +import com.intellij.internal.statistic.eventLog.EventLogGroup +import com.intellij.internal.statistic.service.fus.collectors.CounterUsagesCollector +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.platform.ml.impl.logs.MLEventLoggerProvider.Companion.ML_RECORDER_ID +import com.intellij.platform.ml.impl.tools.registerMLTaskLogging +import com.intellij.util.application +import com.jetbrains.ml.MLUnit +import com.jetbrains.ml.logs.* +import com.jetbrains.ml.logs.schema.BooleanEventField +import com.jetbrains.ml.logs.schema.EventField +import com.jetbrains.ml.logs.schema.EventPair +import com.jetbrains.ml.model.MLModel +import com.jetbrains.ml.session.MLSession +import com.jetbrains.ml.session.MLSessionInfo +import com.jetbrains.ml.tree.LevelFeaturesSchema +import com.jetbrains.ml.tree.MLTree +import com.jetbrains.python.PythonPluginDisposable + + +@Service +private class EventLogGroupHolder { + val group by lazy { + EventLogGroup("pycharm.quickfix.imports", 3, ML_RECORDER_ID).also { + it.registerMLTaskLogging(service().task, + parentDisposable = service()) + } + } +} + +class PyCharmImportsRankingLogs : CounterUsagesCollector() { + override fun getGroup() = service().group +} + +internal enum class LoggingOption { + FULL, + NO_TREE, + SKIP +} + +internal const val RELEASE_LOGGING_PERCENT = 5 + +internal fun getLoggingOption(mlSession: MLSession<*, *>): LoggingOption { + if (application.isEAP || application.isUnitTestMode) { + return LoggingOption.FULL + } + if (mlSession.id % 100 > RELEASE_LOGGING_PERCENT) return LoggingOption.SKIP + return LoggingOption.NO_TREE +} + +object PyCharmImportsRankingLogger : MLSessionLoggerProvider { + private val baseLoggerProvider = EntireSessionLoggerProvider(BooleanEventField("prediction") { "ML model prediction" }) { null } + + override fun createMLSessionLogger( + eventPrefix: String, + taskId: String, + treeAnalysisDeclaration: Map, List>>, + sessionAnalysisDeclaration: List>, + featuresDeclaration: List, + fusEventRegister: FusEventRegister, + ): MLSessionLogger { + val baseSessionLogger = baseLoggerProvider.createMLSessionLogger(eventPrefix, taskId, treeAnalysisDeclaration, sessionAnalysisDeclaration, featuresDeclaration, fusEventRegister) + return object : MLSessionLogger { + override suspend fun logMLSession(sessionInfo: MLSessionInfo, Any>, analysisException: Throwable?, sessionAnalysis: List>, structure: MLTree.ATopNode, Any>?): MLSessionLoggingOutcome { + val loggingOption = checkNotNull(sessionAnalysis.find { it.field == ML_LOGGING_STATE }).data as LoggingOption + return when (loggingOption) { + LoggingOption.FULL -> baseSessionLogger.logMLSession(sessionInfo, analysisException, sessionAnalysis, structure) + LoggingOption.NO_TREE -> { + baseSessionLogger.logMLSession(sessionInfo, analysisException, sessionAnalysis, null) + MLSessionLoggingOutcome.Custom("no tree logged") + } + LoggingOption.SKIP -> MLSessionLoggingOutcome.FilteredOut + } + } + } + } +} diff --git a/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlTask.kt b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlTask.kt new file mode 100644 index 000000000000..54ab019b6de5 --- /dev/null +++ b/python/src/com/jetbrains/python/codeInsight/imports/mlapi/mlTask.kt @@ -0,0 +1,37 @@ +// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license. +package com.jetbrains.python.codeInsight.imports.mlapi + +import com.intellij.codeInsight.inline.completion.ml.MLUnitTyping +import com.intellij.openapi.components.Service +import com.intellij.openapi.components.service +import com.intellij.platform.ml.impl.tools.IJPlatform +import com.jetbrains.ml.MLTaskBuilder +import kotlin.time.Duration.Companion.days + + +internal val LONGEST_POPUP_LIFE_DURATION = 365.days + +@Service +class MLTaskPyCharmImportStatementsRanking { + val task by lazy { + MLTaskBuilder( + taskId = "pycharm_import_statements_ranking", + taskLevels = listOf( + setOf(MLUnitImportCandidatesList, MLUnitTyping), + setOf(MLUnitImportCandidate) + ), + platform = service() + ).buildWithoutMLModel { + + suspendableTreeAnalysis(MLFieldCorrectElement) + suspendableSessionAnalysis(MLFeedbackCorrectElementPosition) + suspendableSessionAnalysis(ML_FEEDBACK_TIME_MS_TO_DISPLAY) + suspendableSessionAnalysis(ML_FEEDBACK_TIME_MS_BEFORE_CLOSED) + runtimeSessionAnalysis(ML_LOGGING_STATE, ML_ENABLED, ML_STATUS_CORRESPONDS_TO_BUCKET) + + logger = PyCharmImportsRankingLogger + + maxAnalysisDuration = LONGEST_POPUP_LIFE_DURATION + } + } +}