Skip to content

Commit c553148

Browse files
DOC-5535 added local copies of source files
1 parent a9ab563 commit c553148

File tree

2 files changed

+517
-0
lines changed

2 files changed

+517
-0
lines changed
Lines changed: 263 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,263 @@
1+
// EXAMPLE: home_vecsets
2+
package io.redis.examples.async;
3+
4+
// STEP_START import
5+
// Lettuce client and query engine classes.
6+
import io.lettuce.core.*;
7+
import io.lettuce.core.api.StatefulRedisConnection;
8+
import io.lettuce.core.api.async.RedisAsyncCommands;
9+
10+
// Standard library classes for data manipulation and
11+
// asynchronous programming.
12+
import java.util.*;
13+
import java.util.concurrent.CompletableFuture;
14+
15+
// DJL classes for model loading and inference.
16+
import ai.djl.huggingface.translator.TextEmbeddingTranslatorFactory;
17+
import ai.djl.inference.Predictor;
18+
import ai.djl.repository.zoo.Criteria;
19+
import ai.djl.training.util.ProgressBar;
20+
// STEP_END
21+
// REMOVE_START
22+
import org.junit.jupiter.api.Test;
23+
import static org.assertj.core.api.Assertions.assertThat;
24+
// REMOVE_END
25+
26+
public class HomeVecSetsExample {
27+
28+
// STEP_START helper_method
29+
public static Double[] convertFloatsToDoubles(float[] input) {
30+
if (input == null) {
31+
return null;
32+
}
33+
34+
Double[] output = new Double[input.length];
35+
36+
for (int i = 0; i < input.length; i++) {
37+
output[i] = Double.valueOf(input[i]);
38+
}
39+
40+
return output;
41+
}
42+
// STEP_END
43+
44+
// REMOVE_START
45+
@Test
46+
// REMOVE_END
47+
public void run() {
48+
// STEP_START model
49+
Predictor<String, float[]> predictor = null;
50+
51+
try {
52+
Criteria<String, float[]> criteria = Criteria.builder().setTypes(String.class, float[].class)
53+
.optModelUrls("djl://ai.djl.huggingface.pytorch/sentence-transformers/all-MiniLM-L6-v2")
54+
.optEngine("PyTorch").optTranslatorFactory(new TextEmbeddingTranslatorFactory())
55+
.optProgress(new ProgressBar()).build();
56+
57+
predictor = criteria.loadModel().newPredictor();
58+
} catch (Exception e) {
59+
// ...
60+
}
61+
// STEP_END
62+
63+
// STEP_START data
64+
final class Person {
65+
final String name;
66+
final int born;
67+
final int died;
68+
final String description;
69+
Person(String name, int born, int died, String description) {
70+
this.name = name; this.born = born; this.died = died; this.description = description;
71+
}
72+
}
73+
74+
List<Person> people = Arrays.asList(
75+
new Person(
76+
"Marie Curie",
77+
1867, 1934,
78+
"Polish-French chemist and physicist. The only person ever to win" +
79+
" two Nobel prizes for two different sciences."
80+
),
81+
new Person(
82+
"Linus Pauling",
83+
1901, 1994,
84+
"American chemist and peace activist. One of only two people to" +
85+
" win two Nobel prizes in different fields (chemistry and peace)."
86+
),
87+
new Person(
88+
"Freddie Mercury",
89+
1946, 1991,
90+
"British musician, best known as the lead singer of the rock band Queen."
91+
),
92+
new Person(
93+
"Marie Fredriksson",
94+
1958, 2019,
95+
"Swedish multi-instrumentalist, mainly known as the lead singer and" +
96+
" keyboardist of the band Roxette."
97+
),
98+
new Person(
99+
"Paul Erdos",
100+
1913, 1996,
101+
"Hungarian mathematician, known for his eccentric personality almost" +
102+
" as much as his contributions to many different fields of mathematics."
103+
),
104+
new Person(
105+
"Maryam Mirzakhani",
106+
1977, 2017,
107+
"Iranian mathematician. The first woman ever to win the Fields medal" +
108+
" for her contributions to mathematics."
109+
),
110+
new Person(
111+
"Masako Natsume",
112+
1957, 1985,
113+
"Japanese actress. She was very famous in Japan but was primarily" +
114+
" known elsewhere in the world for her portrayal of Tripitaka in the" +
115+
" TV series Monkey."
116+
),
117+
new Person(
118+
"Chaim Topol",
119+
1935, 2023,
120+
"Israeli actor and singer, usually credited simply as 'Topol'. He was" +
121+
" best known for his many appearances as Tevye in the musical Fiddler" +
122+
" on the Roof."
123+
)
124+
);
125+
// STEP_END
126+
127+
// STEP_START add_data
128+
RedisClient redisClient = RedisClient.create("redis://localhost:6379");
129+
130+
try (StatefulRedisConnection<String, String> connection = redisClient.connect()) {
131+
RedisAsyncCommands<String, String> asyncCommands = connection.async();
132+
// REMOVE_START
133+
asyncCommands.del("famousPeople").toCompletableFuture().join();
134+
// REMOVE_END
135+
136+
CompletableFuture<?>[] vaddFutures = new CompletableFuture[people.size()];
137+
138+
for (int i = 0; i < people.size(); i++) {
139+
Person person = people.get(i);
140+
Double[] embedding = null;
141+
142+
try {
143+
embedding = convertFloatsToDoubles(predictor.predict(person.description));
144+
} catch (Exception e) {
145+
// ...
146+
}
147+
148+
VAddArgs personArgs = VAddArgs.Builder.attributes(String.format("{\"born\": %d, \"died\": %d}", person.born, person.died));
149+
150+
vaddFutures[i] = asyncCommands.vadd("famousPeople", person.name, personArgs, embedding)
151+
.thenApply(result -> {
152+
System.out.println(result); // >>> true
153+
// REMOVE_START
154+
assertThat(result).isTrue();
155+
// REMOVE_END
156+
return result;
157+
}).toCompletableFuture();
158+
}
159+
160+
CompletableFuture.allOf(vaddFutures).join();
161+
// STEP_END
162+
163+
// STEP_START basic_query
164+
Double[] actorsEmbedding = null;
165+
166+
try {
167+
actorsEmbedding = convertFloatsToDoubles(predictor.predict("actors"));
168+
} catch (Exception e) {
169+
// ...
170+
}
171+
172+
CompletableFuture<List<String>> basicQuery = asyncCommands.vsim("famousPeople", actorsEmbedding)
173+
.thenApply(result -> {
174+
System.out.println(result);
175+
// >>> [Masako Natsume, Chaim Topol, Linus Pauling, Marie Fredriksson, Maryam Mirzakhani, Marie Curie, Freddie Mercury, Paul Erdos]
176+
// REMOVE_START
177+
assertThat(result).containsExactly("Masako Natsume", "Chaim Topol", "Linus Pauling", "Marie Fredriksson", "Maryam Mirzakhani", "Marie Curie", "Freddie Mercury", "Paul Erdos");
178+
// REMOVE_END
179+
return result;
180+
}).toCompletableFuture();
181+
// STEP_END
182+
183+
// STEP_START limited_query
184+
VSimArgs limitArgs = VSimArgs.Builder.count(2L);
185+
186+
CompletableFuture<List<String>> limitedQuery = asyncCommands.vsim("famousPeople", limitArgs, actorsEmbedding)
187+
.thenApply(result -> {
188+
System.out.println(result);
189+
// >>> [Masako Natsume, Chaim Topol]
190+
// REMOVE_START
191+
assertThat(result).containsExactly("Masako Natsume", "Chaim Topol");
192+
// REMOVE_END
193+
return result;
194+
}).toCompletableFuture();
195+
// STEP_END
196+
197+
// STEP_START entertainer_query
198+
Double[] entertainerEmbedding = null;
199+
200+
try {
201+
entertainerEmbedding = convertFloatsToDoubles(predictor.predict("entertainers"));
202+
} catch (Exception e) {
203+
// ...
204+
}
205+
206+
CompletableFuture<List<String>> entertainerQuery = asyncCommands.vsim("famousPeople", entertainerEmbedding)
207+
.thenApply(result -> {
208+
System.out.println(result);
209+
// >>> [Freddie Mercury, Chaim Topol, Linus Pauling, Marie Fredriksson, Masako Natsume, Paul Erdos, Maryam Mirzakhani, Marie Curie]
210+
// REMOVE_START
211+
assertThat(result).containsExactly("Freddie Mercury", "Chaim Topol", "Linus Pauling", "Marie Fredriksson", "Masako Natsume", "Paul Erdos", "Maryam Mirzakhani", "Marie Curie");
212+
// REMOVE_END
213+
return result;
214+
}).toCompletableFuture();
215+
// STEP_END
216+
217+
Double[] scienceEmbedding = null;
218+
219+
try {
220+
scienceEmbedding = convertFloatsToDoubles(predictor.predict("science"));
221+
} catch (Exception e) {
222+
// ...
223+
}
224+
225+
CompletableFuture<List<String>> scienceQuery = asyncCommands.vsim("famousPeople", scienceEmbedding)
226+
.thenApply(result -> {
227+
System.out.println(result);
228+
// >>> [Marie Curie, Linus Pauling, Maryam Mirzakhani, Paul Erdos, Marie Fredriksson, Freddie Mercury, Masako Natsume, Chaim Topol]
229+
// REMOVE_START
230+
assertThat(result).containsExactly("Marie Curie", "Linus Pauling", "Maryam Mirzakhani", "Paul Erdos", "Marie Fredriksson", "Freddie Mercury", "Masako Natsume", "Chaim Topol");
231+
// REMOVE_END
232+
return result;
233+
}).toCompletableFuture();
234+
235+
// STEP_START filtered_query
236+
Double[] science2000Embedding = null;
237+
238+
try {
239+
science2000Embedding = convertFloatsToDoubles(predictor.predict("science"));
240+
} catch (Exception e) {
241+
// ...
242+
}
243+
244+
VSimArgs filterArgs = VSimArgs.Builder.filter(".died < 2000");
245+
246+
CompletableFuture<List<String>> filteredQuery = asyncCommands.vsim("famousPeople", filterArgs, science2000Embedding)
247+
.thenApply(result -> {
248+
System.out.println(result);
249+
// >>> [Marie Curie, Linus Pauling, Paul Erdos, Freddie Mercury, Masako Natsume]
250+
// REMOVE_START
251+
assertThat(result).containsExactly("Marie Curie", "Linus Pauling", "Paul Erdos", "Freddie Mercury", "Masako Natsume");
252+
// REMOVE_END
253+
return result;
254+
}).toCompletableFuture();
255+
// STEP_END
256+
257+
CompletableFuture.allOf(basicQuery, limitedQuery, entertainerQuery, scienceQuery, filteredQuery).join();
258+
} finally {
259+
redisClient.shutdown();
260+
}
261+
}
262+
}
263+

0 commit comments

Comments
 (0)