From 1c369e3a4e1c9a340fd08762a5e9f3ef387052d4 Mon Sep 17 00:00:00 2001
From: Kirill Golikov <kirill.golikov@abbyy.com>
Date: Mon, 9 Sep 2024 20:21:40 +0200
Subject: [PATCH] [NeoML] CTiedEmbeddingsLayer extend functional interface
 (#1107)

Signed-off-by: Kirill Golikov <kirill.golikov@abbyy.com>
---
 NeoML/include/NeoML/Dnn/DnnLambdaHolder.h     | 57 +++++++++++--------
 .../NeoML/Dnn/Layers/TiedEmbeddingsLayer.h    |  3 +-
 NeoML/src/Dnn/Layers/TiedEmbeddingsLayer.cpp  | 14 +++--
 3 files changed, 45 insertions(+), 29 deletions(-)

diff --git a/NeoML/include/NeoML/Dnn/DnnLambdaHolder.h b/NeoML/include/NeoML/Dnn/DnnLambdaHolder.h
index 510554700..a627001d0 100644
--- a/NeoML/include/NeoML/Dnn/DnnLambdaHolder.h
+++ b/NeoML/include/NeoML/Dnn/DnnLambdaHolder.h
@@ -1,4 +1,4 @@
-/* Copyright © 2017-2020 ABBYY Production LLC
+/* Copyright © 2017-2024 ABBYY
 
 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
@@ -15,58 +15,70 @@ limitations under the License.
 
 #pragma once
 
+#include <type_traits>
 #include <NeoML/NeoMLDefs.h>
 
 namespace NeoML {
 
-//////////////////////////////////////////////////////////////////////////////////////////
-
 // Simple analog for std::function() that does not use std::allocator
 
+namespace details {
+
 template<typename T>
-class CLambdaHolderBase {};
+class ILambdaHolderBase;
 
 // Base class for lambda holder. This interface hide actual lambda type.
 template<typename Out, typename ...In>
-class CLambdaHolderBase<Out( In... )> : public virtual IObject {
+class ILambdaHolderBase<Out( In... )> : public virtual IObject {
 public:
 	// Executes lambda.
 	virtual Out Execute( In... arguments ) = 0;
-	// Copies lambda.
-	virtual CPtr<CLambdaHolderBase<Out( In... )>> Copy() = 0;
 };
 
+//---------------------------------------------------------------------------------------------------------------------
+
 template<typename T, typename U>
-class CLambdaHolder {};
+class CLambdaHolder;
 
 // Lambda holder implementation.
-template<typename T, typename Out, typename ...In>
-class CLambdaHolder<T, Out( In... )> : public CLambdaHolderBase<Out( In... )> {
+template<typename F, typename Out, typename ...In>
+class CLambdaHolder<F, Out( In... )> : public ILambdaHolderBase<Out( In... )> {
 public:
-	CLambdaHolder( T _lambda ) : lambda( _lambda ) {}
+	CLambdaHolder( F&& func ) : lambda( std::move( func ) ) {}
+	CLambdaHolder( const F& func ) : lambda( func ) {}
 
 	Out Execute( In... in ) override
 		{ return lambda( in... ); }
 
-	CPtr<CLambdaHolderBase<Out( In... )>> Copy() override
-		{ return new CLambdaHolder<T, Out( In... )>( lambda ); }
-
 private:
-	T lambda;
+	F lambda;
 };
 
+} // namespace details
+
+//---------------------------------------------------------------------------------------------------------------------
+
 template<typename T>
-class CLambda {};
+class CLambda;
 
 // Type that captures lambda.
 template<typename Out, typename ...In>
 class CLambda<Out( In... )> {
 public:
-	CLambda() {}
-	template<class T>
-	CLambda( const T& t ) : lambda( new CLambdaHolder<T, Out( In... )>( t ) ) {}
-	CLambda( const CLambda& other ) :
-		lambda( other.lambda != 0 ? other.lambda->Copy() : nullptr ) {}
+	CLambda() = default;
+	// Be copied and moved by default, because it stores the shared pointer
+
+	// Convert from a function, except itself type
+	// By coping
+	template<class F,
+		typename std::enable_if<!std::is_same<CLambda, typename std::decay<F>::type>::value, int>::type = 0>
+	CLambda( const F& function ) :
+		lambda( new details::CLambdaHolder<F, Out( In... )>( function ) ) {}
+	// By moving
+	template<class F,
+		typename std::enable_if<!std::is_same<CLambda, typename std::decay<F>::type>::value, int>::type = 0>
+	CLambda( F&& function ) :
+		lambda( new details::CLambdaHolder<F, Out( In... )>( std::move( function ) ) ) {}
 
 	Out operator()( In... in )
 	{
@@ -78,8 +90,7 @@ class CLambda<Out( In... )> {
 	bool IsEmpty() const { return lambda == nullptr; }
 
 private:
-	CPtr<CLambdaHolderBase<Out( In... )>> lambda;
+	CPtr<details::ILambdaHolderBase<Out( In... )>> lambda;
 };
 
-//////////////////////////////////////////////////////////////////////////////////////////
 } // namespace NeoML
diff --git a/NeoML/include/NeoML/Dnn/Layers/TiedEmbeddingsLayer.h b/NeoML/include/NeoML/Dnn/Layers/TiedEmbeddingsLayer.h
index b922b758f..0d478b1b2 100644
--- a/NeoML/include/NeoML/Dnn/Layers/TiedEmbeddingsLayer.h
+++ b/NeoML/include/NeoML/Dnn/Layers/TiedEmbeddingsLayer.h
@@ -71,6 +71,7 @@ class NEOML_API CTiedEmbeddingsLayer : public CBaseLayer {
 };
 
 // Tied embeddings.
-NEOML_API CLayerWrapper<CTiedEmbeddingsLayer> TiedEmbeddings( const char* name, int channel );
+NEOML_API CLayerWrapper<CTiedEmbeddingsLayer> TiedEmbeddings( const char* name, int channel,
+	CArray<CString>&& embeddingPath = {} );
 
 } // namespace NeoML
diff --git a/NeoML/src/Dnn/Layers/TiedEmbeddingsLayer.cpp b/NeoML/src/Dnn/Layers/TiedEmbeddingsLayer.cpp
index 7c3c847d4..b4ed7f431 100644
--- a/NeoML/src/Dnn/Layers/TiedEmbeddingsLayer.cpp
+++ b/NeoML/src/Dnn/Layers/TiedEmbeddingsLayer.cpp
@@ -156,12 +156,16 @@ const CMultichannelLookupLayer* CTiedEmbeddingsLayer::getLookUpLayer() const
 	return embeddingsLayer;
 }
 
-CLayerWrapper<CTiedEmbeddingsLayer> TiedEmbeddings( const char* name, int channel )
+CLayerWrapper<CTiedEmbeddingsLayer> TiedEmbeddings( const char* name, int channel, CArray<CString>&& embeddingPath )
 {
-	return CLayerWrapper<CTiedEmbeddingsLayer>( "TiedEmbeddings", [=]( CTiedEmbeddingsLayer* result ) {
-		result->SetEmbeddingsLayerName( name );
-		result->SetChannelIndex( channel );
-	} );
+	return CLayerWrapper<CTiedEmbeddingsLayer>( "TiedEmbeddings",
+		[=, path=std::move( embeddingPath )]( CTiedEmbeddingsLayer* result )
+		{
+			result->SetEmbeddingsLayerName( name );
+			result->SetChannelIndex( channel );
+			result->SetEmbeddingsLayerPath( path );
+		}
+	);
 }
 
 } // namespace NeoML