索引和查询向量

dotnet new console -n VecQueryExample

dotnet add package NRedisStack

dotnet add package Microsoft.ML

// Redis connection and Query Engine.
using NRedisStack.RedisStackCommands;
using StackExchange.Redis;
using NRedisStack.Search;
using static NRedisStack.Search.Schema;
using NRedisStack.Search.Literals.Enums;

// Text embeddings.
using Microsoft.ML;
using Microsoft.ML.Transforms.Text;

// Azure embeddings.
using Azure;
using Azure.AI.OpenAI;

static PredictionEngine<TextData, TransformedTextData> GetPredictionEngine(){
    // Create a new ML context, for ML.NET operations. It can be used for
    // exception tracking and logging, as well as the source of randomness.
    var mlContext = new MLContext();

    // Create an empty list as the dataset
    var emptySamples = new List<TextData>();

    // Convert sample list to an empty IDataView.
    var emptyDataView = mlContext.Data.LoadFromEnumerable(emptySamples);

    // A pipeline for converting text into a 150-dimension embedding vector
    var textPipeline = mlContext.Transforms.Text.NormalizeText("Text")
        .Append(mlContext.Transforms.Text.TokenizeIntoWords("Tokens",
            "Text"))
        .Append(mlContext.Transforms.Text.ApplyWordEmbedding("Features",
            "Tokens", WordEmbeddingEstimator.PretrainedModelKind
            .SentimentSpecificWordEmbedding));

    // Fit to data.
    var textTransformer = textPipeline.Fit(emptyDataView);

    // Create the prediction engine to get the embedding vector from the input text/string.
    var predictionEngine = mlContext.Model.CreatePredictionEngine<TextData,
        TransformedTextData>(textTransformer);

    return predictionEngine;
}

static byte[] GetEmbedding(
    PredictionEngine<TextData, TransformedTextData> model, string sentence
)
{
    // Call the prediction API to convert the text into embedding vector.
    var data = new TextData()
    {
        Text = sentence
    };

    var prediction = model.Predict(data);

    // Convert prediction.Features to a binary blob
    float[] floatArray = Array.ConvertAll(prediction.Features, x => (float)x);
    byte[] byteArray = new byte[floatArray.Length * sizeof(float)];
    Buffer.BlockCopy(floatArray, 0, byteArray, 0, byteArray.Length);

    return byteArray;
}

private static byte[] GetEmbeddingFromAzure(string sentence){
	Uri oaiEndpoint = new ("your-azure-openai-endpoint”);
	string oaiKey = "your-openai-key";

	AzureKeyCredential credentials = new (oaiKey);
	OpenAIClient openAIClient = new (oaiEndpoint, credentials);

	EmbeddingsOptions embeddingOptions = new() {
    	     DeploymentName = "your-deployment-name",
    	     Input = { sentence },
	};

	// Generate the vector embedding
	var returnValue = openAIClient.GetEmbeddings(embeddingOptions);

	// Convert the array of floats to binary blob
	float[] floatArray = Array.ConvertAll(returnValue.Value.Data[0].Embedding.ToArray(), x => (float)x);
	byte[] byteArray = new byte[floatArray.Length * sizeof(float)];
	Buffer.BlockCopy(floatArray, 0, byteArray, 0, byteArray.Length);
	return byteArray;
}

var muxer = ConnectionMultiplexer.Connect("localhost:6379");
var db = muxer.GetDatabase();

try { db.FT().DropIndex("vector_idx");} catch {}

var schema = new Schema()
    .AddTextField(new FieldName("content", "content"))
    .AddTagField(new FieldName("genre", "genre"))
    .AddVectorField("embedding", VectorField.VectorAlgo.HNSW,
        new Dictionary<string, object>()
        {
            ["TYPE"] = "FLOAT32",
            ["DIM"] = "150",
            ["DISTANCE_METRIC"] = "L2"
        }
    );

db.FT().Create(
    "vector_idx",
    new FTCreateParams()
        .On(IndexDataType.HASH)
        .Prefix("doc:"),
    schema
);

var predEngine = GetPredictionEngine();

var sentence1 = "That is a very happy person";

HashEntry[] doc1 = {
    new("content", sentence1),
    new("genre", "persons"),
    new("embedding", GetEmbedding(predEngine, sentence1))
};

db.HashSet("doc:1", doc1);

var sentence2 = "That is a happy dog";

HashEntry[] doc2 = {
    new("content", sentence2),
    new("genre", "pets"),
    new("embedding", GetEmbedding(predEngine, sentence2))
};

db.HashSet("doc:2", doc2);

var sentence3 = "Today is a sunny day";

HashEntry[] doc3 = {
    new("content", sentence3),
    new("genre", "weather"),
    new("embedding", GetEmbedding(predEngine, sentence3))
};

db.HashSet("doc:3", doc3);

var res = db.FT().Search("vector_idx",
    new Query("*=>[KNN 3 @embedding $query_vec AS score]")
    .AddParam("query_vec", GetEmbedding(predEngine, "That is a happy person"))
    .ReturnFields(
        new FieldName("content", "content"),
        new FieldName("score", "score")
    )
    .SetSortBy("score")
    .Dialect(2));

foreach (var doc in res.Documents) {
    var props = doc.GetProperties();
    var propText = string.Join(
        ", ",
        props.Select(p => $"{p.Key}: '{p.Value}'")
    );

    Console.WriteLine(
        $"ID: {doc.Id}, Properties: [\n  {propText}\n]"
    );
}

class TextData
{
    public string Text { get; set; }
}

class TransformedTextData : TextData
{
    public float[] Features { get; set; }
}

var jsonSchema = new Schema()
    .AddTextField(new FieldName("$.content", "content"))
    .AddTagField(new FieldName("$.genre", "genre"))
    .AddVectorField(
        new FieldName("$.embedding", "embedding"),
        VectorField.VectorAlgo.HNSW,
        new Dictionary<string, object>()
        {
            ["TYPE"] = "FLOAT32",
            ["DIM"] = "150",
            ["DISTANCE_METRIC"] = "L2"
        }
    );


db.FT().Create(
    "vector_json_idx",
    new FTCreateParams()
        .On(IndexDataType.JSON)
        .Prefix("jdoc:"),
    jsonSchema
);

static float[] GetFloatEmbedding(
    PredictionEngine<TextData, TransformedTextData> model, string sentence
)
{
    // Call the prediction API to convert the text into embedding vector.
    var data = new TextData()
    {
        Text = sentence
    };

    var prediction = model.Predict(data);

    float[] floatArray = Array.ConvertAll(prediction.Features, x => (float)x);
    return floatArray;
}

var jSentence1 = "That is a very happy person";

var jdoc1 = new {
    content = jSentence1,
    genre = "persons",
    embedding = GetFloatEmbedding(predEngine, jSentence1),
};

db.JSON().Set("jdoc:1", "$", jdoc1);

var jSentence2 = "That is a happy dog";

var jdoc2 = new {
    content = jSentence2,
    genre = "pets",
    embedding = GetFloatEmbedding(predEngine, jSentence2),
};

db.JSON().Set("jdoc:2", "$", jdoc2);

var jSentence3 = "Today is a sunny day";

var jdoc3 = new {
    content = jSentence3,
    genre = "weather",
    embedding = GetFloatEmbedding(predEngine, jSentence3),
};

db.JSON().Set("jdoc:3", "$", jdoc3);

var jRes = db.FT().Search("vector_json_idx",
    new Query("*=>[KNN 3 @embedding $query_vec AS score]")
    .AddParam("query_vec", GetEmbedding(predEngine, "That is a happy person"))
    .ReturnFields(
        new FieldName("$.content", "content"),
        new FieldName("$.score", "score")
    )
    .SetSortBy("score")
    .Dialect(2));

foreach (var doc in jRes.Documents) {
    var props = doc.GetProperties();
    var propText = string.Join(
        ", ",
        props.Select(p => $"{p.Key}: '{p.Value}'")
    );

    Console.WriteLine(
        $"ID: {doc.Id}, Properties: [\n  {propText}\n]"
    );
}