#4292 - ollama-based recommender

- Added few-shot example - Fix handling of JSON responses consisting of numbers
inception-project · Nov 28, 2023 · efb78fa · efb78fa
1 parent 5162dae
commit efb78fa
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 9 deletions.
diff --git a/...armstadt/ukp/inception/recommendation/imls/ollama/response/MentionsFromJsonExtractor.java b/...armstadt/ukp/inception/recommendation/imls/ollama/response/MentionsFromJsonExtractor.java
@@ -113,7 +113,7 @@ List<Pair<String, String>> extractMentionFromJson(String aResponse)
                 var fieldEntry = fieldIterator.next();
                 if (fieldEntry.getValue().isArray()) {
                     for (var item : fieldEntry.getValue()) {
-                        if (item.isTextual()) {
+                        if (item.isTextual() || item.isNumber() || item.isBoolean()) {
                             // Looks like this
                             // "Person": ["John"],
                             // "Location": ["diner", "Starbucks"]

diff --git a/...a/src/main/resources/de/tudarmstadt/ukp/inception/recommendation/imls/ollama/presets.yaml b/...a/src/main/resources/de/tudarmstadt/ukp/inception/recommendation/imls/ollama/presets.yaml
@@ -8,6 +8,30 @@
 
     {{ text }}</s>
 
+- name: Extract named entities from sentenes (dynamic few-shot)
+  promptingMode: per-sentence
+  format: json
+  extractionMode: mentions-from-json
+  prompt: |-
+    Identify all politicians in the following text and return them as JSON.
+
+    {% for example in examples %}
+    Text:
+    '''
+    {{ example.getText() }}
+    '''
+
+    Response:
+    {{ example.getLabelledMentions() | tojson }}
+    {% endfor %}
+
+    Text:
+    '''
+    {{ text }}
+    '''
+
+    Response:
+
 - name: Summarize each sentence in a word
   promptingMode: per-sentence
   extractionMode: response-as-label

diff --git a/...t/java/de/tudarmstadt/ukp/inception/recommendation/imls/ollama/OllamaRecommenderTest.java b/...t/java/de/tudarmstadt/ukp/inception/recommendation/imls/ollama/OllamaRecommenderTest.java
@@ -111,8 +111,10 @@ void testPerDocumentUsingMentionsFromJsonList_Numbers() throws Exception
 
         var traits = new OllamaRecommenderTraits();
         traits.setModel("mistral");
-        traits.setPrompt(
-                "Identify all even numbers in the following list and return them as JSON.\n\n{{ text }}");
+        traits.setPrompt("""
+                Identify all even numbers in the following list and return them as JSON.
+
+                {{ text }}""");
         traits.setFormat(JSON);
         traits.setPromptingMode(PER_DOCUMENT);
         traits.setExtractionMode(MENTIONS_FROM_JSON);

diff --git a/...tadt/ukp/inception/recommendation/imls/ollama/response/MentionsFromJsonExtractorTest.java b/...tadt/ukp/inception/recommendation/imls/ollama/response/MentionsFromJsonExtractorTest.java
@@ -28,7 +28,22 @@ class MentionsFromJsonExtractorTest
     private MentionsFromJsonExtractor sut = new MentionsFromJsonExtractor();
 
     @Test
-    void testExtractMentionFromJson_variant1()
+    void testExtractMentionFromJson_categorizedNumbers()
+    {
+        var json = """
+                {
+                    "even_numbers": [2, 4, 6]
+                }
+                """;
+        assertThat(sut.extractMentionFromJson(json)) //
+                .containsExactly( //
+                        Pair.of("2", "even_numbers"), //
+                        Pair.of("4", "even_numbers"), //
+                        Pair.of("6", "even_numbers"));
+    }
+
+    @Test
+    void testExtractMentionFromJson_categorizedStrings()
     {
         var json = """
                 {
@@ -44,7 +59,7 @@ void testExtractMentionFromJson_variant1()
     }
 
     @Test
-    void testExtractMentionFromJson_variant2()
+    void testExtractMentionFromJson_categorizedObjects()
     {
         var json = """
                 {
@@ -63,7 +78,7 @@ void testExtractMentionFromJson_variant2()
     }
 
     @Test
-    void testExtractMentionFromJson_variant3()
+    void testExtractMentionFromJson_namedObjects()
     {
         var json = """
                 {
@@ -80,7 +95,7 @@ void testExtractMentionFromJson_variant3()
     }
 
     @Test
-    void testExtractMentionFromJson_variant4()
+    void testExtractMentionFromJson_keyValue()
     {
         var json = """
                 {
@@ -98,9 +113,9 @@ void testExtractMentionFromJson_variant4()
                         Pair.of("Don Horny", "politician"));
     }
 
-    @Disabled("Cannot really tell this one apart from variant 4")
+    @Disabled("Cannot really tell this one apart from keyValue")
     @Test
-    void testExtractMentionFromJson_variant5()
+    void testExtractMentionFromJson_valueKey()
     {
         // We assume that the first item is the most relevant one (the
         // mention) so we do not get a bad mention in cases like this: