diff --git a/frontend/components/core/ReTableInfo.vue b/frontend/components/core/ReTableInfo.vue
index fa14679455..d4e8f09ea3 100644
--- a/frontend/components/core/ReTableInfo.vue
+++ b/frontend/components/core/ReTableInfo.vue
@@ -111,8 +111,16 @@
                       {{ itemValue(item, column) | percent }}
                     </span>
                     <span v-else-if="column.type === 'array'">
-                      <p v-for="(arrayItem, index) in itemValue(item, column)" :key="index">
-                        {{arrayItem}}{{index + 1 === itemValue(item, column).length ? '' : ','}}
+                      <p
+                        v-for="(arrayItem, index) in itemValue(item, column)"
+                        :key="index"
+                      >
+                        {{ arrayItem
+                        }}{{
+                          index + 1 === itemValue(item, column).length
+                            ? ""
+                            : ","
+                        }}
                       </p>
                     </span>
                     <span v-else-if="column.type === 'object'">
diff --git a/frontend/components/text-classifier/results/RecordExplain.vue b/frontend/components/text-classifier/results/RecordExplain.vue
index 742fa08b56..cb8f798490 100644
--- a/frontend/components/text-classifier/results/RecordExplain.vue
+++ b/frontend/components/text-classifier/results/RecordExplain.vue
@@ -33,19 +33,18 @@
 <script>
 export default {
   props: {
-    queryText: {
-      type: String,
-      default: undefined,
+    record: {
+      type: Object,
+      required: true,
     },
     explain: {
       type: Array,
-    },
-    predicted: {
-      type: String,
-      default: undefined,
-    },
+    }
   },
   computed: {
+    predicted() {
+      return this.record.predicted;
+    },
     explainFormatted() {
       // TODO ALLOW FOR MULTI LABEL
       return this.explain.map((token) => {
@@ -58,8 +57,8 @@ export default {
           percent = Math.round(Math.log10(percent) ** p * s);
         }
         return {
-          text: this.queryText
-            ? this.$highlightSearch(this.queryText, token.token)
+          text: this.record.search_keywords
+            ? this.$highlightKeywords(token.token, this.record.search_keywords)
             : token.token,
           percent: percent.toString(),
           grad,
diff --git a/frontend/components/text-classifier/results/RecordInputs.vue b/frontend/components/text-classifier/results/RecordInputs.vue
index a9075d6c39..a0a162fb79 100644
--- a/frontend/components/text-classifier/results/RecordInputs.vue
+++ b/frontend/components/text-classifier/results/RecordInputs.vue
@@ -26,11 +26,10 @@
           <span class="record__key">{{ index }}:</span>
           <LazyRecordExplain
             v-if="explanation"
-            :predicted="predicted"
-            :query-text="queryText"
+            :record="record"
             :explain="explanation[index]"
           />
-          <LazyRecordString v-else :query-text="queryText" :text="text" />
+          <LazyRecordString v-else :record="record" :text="text" />
         </span>
       </span>
     </div>
@@ -47,27 +46,22 @@
 <script>
 export default {
   props: {
-    data: {
+    record: {
       type: Object,
-      required: true,
-    },
-    queryText: {
-      type: String,
-    },
-    predicted: {
-      type: String,
-      default: undefined,
-    },
-    explanation: {
-      type: Object,
-      default: () => undefined,
-    },
+      required: true
+    }
   },
   data: () => ({
     showFullRecord: false,
     scrollHeight: undefined,
   }),
   computed: {
+    data() {
+      return this.record.inputs;
+    },
+    explanation() {
+      return this.record.explanation;
+    },
     visibleRecordHeight() {
       return this.$mq === "lg" ? 468 : 174;
     },
diff --git a/frontend/components/text-classifier/results/RecordString.vue b/frontend/components/text-classifier/results/RecordString.vue
index 761bb085b5..e1ef1473b5 100644
--- a/frontend/components/text-classifier/results/RecordString.vue
+++ b/frontend/components/text-classifier/results/RecordString.vue
@@ -21,7 +21,7 @@
       <div v-for="item in text" :key="item.index">
         <span
           class="record__content"
-          v-html="$highlightSearch(queryText, item)"
+          v-html="$highlightKeywords(item, keywords)"
         >
         </span>
       </div>
@@ -30,22 +30,28 @@
     <span
       v-else
       class="record__content"
-      v-html="$highlightSearch(queryText, text)"
+      v-html="$highlightKeywords(text, keywords)"
     >
     </span>
   </div>
 </template>
+
 <script>
 export default {
   props: {
+    record: {
+      type: Object,
+      required: true,
+    },
     text: {
       type: [String, Array],
       required: true,
     },
-    queryText: {
-      type: String,
-      default: undefined,
-    },
+  },
+  computed: {
+    keywords() {
+      return this.record.search_keywords
+    }
   },
   methods: {
     isList(record) {
@@ -54,6 +60,7 @@ export default {
   },
 };
 </script>
+
 <style lang="scss" scoped>
 .record {
   &__content {
diff --git a/frontend/components/text-classifier/results/RecordTextClassification.vue b/frontend/components/text-classifier/results/RecordTextClassification.vue
index 6ce518b704..27dfe8fcbf 100644
--- a/frontend/components/text-classifier/results/RecordTextClassification.vue
+++ b/frontend/components/text-classifier/results/RecordTextClassification.vue
@@ -20,12 +20,7 @@
     <!-- annotation labels and prediction status -->
     <div class="record--left">
       <!-- record text -->
-      <RecordInputs
-        :predicted="record.predicted"
-        :data="record.inputs"
-        :explanation="record.explanation"
-        :query-text="dataset.query.text"
-      />
+      <RecordInputs :record="record" />
       <ClassifierAnnotationArea
         v-if="annotationEnabled"
         :dataset="dataset"
diff --git a/frontend/components/text2text/results/RecordStringText2Text.vue b/frontend/components/text2text/results/RecordStringText2Text.vue
index 60a0fc543f..7c14e7e522 100644
--- a/frontend/components/text2text/results/RecordStringText2Text.vue
+++ b/frontend/components/text2text/results/RecordStringText2Text.vue
@@ -20,7 +20,7 @@
     ref="list"
     :class="showFullRecord ? 'record__expanded' : 'record__collapsed'"
   >
-    <span class="record__content" v-html="$highlightSearch(queryText, text)">
+    <span class="record__content" v-html="$highlightKeywords(text, keywords)">
     </span>
     <a
       href="#"
@@ -34,20 +34,22 @@
 <script>
 export default {
   props: {
-    text: {
-      type: [String, Array],
-      required: true,
-    },
-    queryText: {
-      type: String,
-      default: undefined,
-    },
+    record: {
+      type: Object,
+      required: true
+    }
   },
   data: () => ({
     showFullRecord: false,
     scrollHeight: undefined,
   }),
   computed: {
+    text() {
+      return this.record.text;
+    },
+    keywords() {
+      return this.record.search_keywords;
+    },
     visibleRecordHeight() {
       return this.$mq === "lg" ? 570 : 260;
     },
diff --git a/frontend/components/text2text/results/RecordText2Text.vue b/frontend/components/text2text/results/RecordText2Text.vue
index b5bc84a5ce..170f8ec410 100644
--- a/frontend/components/text2text/results/RecordText2Text.vue
+++ b/frontend/components/text2text/results/RecordText2Text.vue
@@ -18,10 +18,7 @@
 <template>
   <div class="record">
     <div class="record--left record__item">
-      <record-string-text-2-text
-        :query-text="dataset.query.text"
-        :text="record.text"
-      />
+      <record-string-text-2-text :record="record" />
       <div>
         <text-2-text-list
           ref="list"
diff --git a/frontend/components/token-classifier/results/EntityHighlight.vue b/frontend/components/token-classifier/results/EntityHighlight.vue
index 5f852487f7..b7f7a27610 100755
--- a/frontend/components/token-classifier/results/EntityHighlight.vue
+++ b/frontend/components/token-classifier/results/EntityHighlight.vue
@@ -19,12 +19,7 @@
   <span
     @mouseenter="showTooltip = true"
     @mouseleave="showTooltip = false"
-    :class="[
-      'highlight',
-      span.origin,
-      // isText ? '' : 'highlight--block',
-      annotationEnabled ? 'editable' : null,
-    ]"
+    :class="['highlight', span.origin, annotationEnabled ? 'editable' : null]"
     ><span
       v-for="(token, i) in span.tokens"
       :key="i"
@@ -36,11 +31,7 @@
       ]"
       @click="openTagSelector"
       @dblclick="removeEntity"
-      v-html="
-        `${$highlightSearch(dataset.query.text, token.text)}${
-          token.hasSpaceAfter && i + 1 !== span.tokens.length ? ' ' : ''
-        }`
-      "
+      v-html="visualizeToken(token, i)"
     ></span>
     <svgicon
       class="remove-button"
@@ -53,6 +44,7 @@
     <lazy-text-span-tooltip v-if="showTooltip" :span="span" />
   </span>
 </template>
+
 <script>
 import "assets/icons/cross";
 
@@ -69,6 +61,10 @@ export default {
       type: Object,
       required: true,
     },
+    record: {
+      type: Object,
+      required: true,
+    }
   },
   data: () => {
     return {
@@ -79,9 +75,6 @@ export default {
     };
   },
   computed: {
-    // isText() {
-    //   return this.text.replace(/\s/g, "").length;
-    // },
     annotationEnabled() {
       return this.dataset.viewSettings.viewMode === "annotate";
     },
@@ -107,19 +100,26 @@ export default {
         }, this.singleClickDelay);
       }
     },
+    visualizeToken(token, i) {
+      let text = token.highlighted
+        ? this.$htmlHighlightText(token.text)
+        : this.$htmlText(token.text);
+      return `${text}${token.hasSpaceAfter && i + 1 !== this.span.tokens.length ? " " : ""}`;
+    },
   },
 };
 </script>
+
 <style lang="scss" scoped>
 .whitespace {
   margin-right: 3.5px;
 }
+
 .highlight {
   @include font-size(0);
   line-height: 1em;
   position: relative;
-  cursor: default;
-  // display: inline-flex;
+  cursor: default; // display: inline-flex;
   border-radius: 2px;
   padding: 0;
   &.editable {
@@ -146,6 +146,7 @@ export default {
     z-index: 5;
   }
 }
+
 .remove-button {
   opacity: 0;
   z-index: -1;
diff --git a/frontend/components/token-classifier/results/RecordTokenClassification.vue b/frontend/components/token-classifier/results/RecordTokenClassification.vue
index 4e1cbea591..77a418aead 100755
--- a/frontend/components/token-classifier/results/RecordTokenClassification.vue
+++ b/frontend/components/token-classifier/results/RecordTokenClassification.vue
@@ -21,11 +21,13 @@
       <record-token-classification-annotation
         :dataset="dataset"
         :record="record"
+        :visualTokens="visualTokens"
         v-if="annotationEnabled"
       />
       <record-token-classification-exploration
         :dataset="dataset"
         :record="record"
+        :visualTokens="visualTokens"
         v-else
       />
     </div>
@@ -33,6 +35,8 @@
 </template>
 
 <script>
+import { indexOf, length } from "stringz";
+
 export default {
   props: {
     dataset: {
@@ -48,6 +52,44 @@ export default {
     annotationEnabled() {
       return this.dataset.viewSettings.viewMode === "annotate";
     },
+    visualTokens() {
+      // This is used for both, annotation ad exploration components
+      const recordHasEmoji = this.record.text.containsEmoji;
+      const searchKeywordsSpans = this.$keywordsSpans(
+        this.record.text,
+        this.record.search_keywords
+      );
+
+      const { visualTokens } = this.record.tokens.reduce(
+        ({ visualTokens, startPosition }, token) => {
+          const start = recordHasEmoji
+            ? indexOf(this.record.text, token, startPosition)
+            : this.record.text.indexOf(token, startPosition);
+          const end = start + (recordHasEmoji ? length(token) : token.length);
+          const hasSpaceAfter = this.record.text.slice(end, end + 1) === " ";
+
+          let highlighted = false;
+          for (let highlight of searchKeywordsSpans) {
+            if (highlight.start <= start && highlight.end >= end) {
+              highlighted = true;
+              break;
+            }
+          }
+          return {
+            visualTokens: [
+              ...visualTokens,
+              { start, end, highlighted, text: token, hasSpaceAfter },
+            ],
+            startPosition: end,
+          };
+        },
+        {
+          visualTokens: [],
+          startPosition: 0,
+        }
+      );
+      return visualTokens;
+    },
   },
 };
 </script>
diff --git a/frontend/components/token-classifier/results/RecordTokenClassificationAnnotation.vue b/frontend/components/token-classifier/results/RecordTokenClassificationAnnotation.vue
index 85f3e0187b..03ce83f10e 100755
--- a/frontend/components/token-classifier/results/RecordTokenClassificationAnnotation.vue
+++ b/frontend/components/token-classifier/results/RecordTokenClassificationAnnotation.vue
@@ -22,18 +22,20 @@
         v-if="record.prediction"
         v-once
         key="prediction"
-        :dataset="dataset"
         origin="prediction"
-        :record="record"
         class="prediction"
+        :dataset="dataset"
+        :record="record"
+        :visualTokens="visualTokens"
         :entities="getEntitiesByOrigin('prediction')"
       />
       <text-spans
         key="annotation"
-        :dataset="dataset"
         origin="annotation"
-        :record="record"
         class="annotation"
+        :dataset="dataset"
+        :record="record"
+        :visualTokens="visualTokens"
         :entities="getEntitiesByOrigin('annotation')"
       />
     </div>
@@ -64,6 +66,10 @@ export default {
       type: Object,
       required: true,
     },
+    visualTokens: {
+      type: Array,
+      required: true,
+    },
   },
   methods: {
     ...mapActions({
diff --git a/frontend/components/token-classifier/results/RecordTokenClassificationExploration.vue b/frontend/components/token-classifier/results/RecordTokenClassificationExploration.vue
index 431f664870..e067eff659 100755
--- a/frontend/components/token-classifier/results/RecordTokenClassificationExploration.vue
+++ b/frontend/components/token-classifier/results/RecordTokenClassificationExploration.vue
@@ -21,18 +21,20 @@
       v-if="record.prediction"
       v-once
       key="prediction"
-      :dataset="dataset"
       origin="prediction"
-      :record="record"
       class="prediction"
+      :dataset="dataset"
+      :record="record"
+      :visualTokens="visualTokens"
       :entities="getEntitiesByOrigin('prediction')"
     />
     <text-spans
       key="annotation"
-      :dataset="dataset"
       origin="annotation"
-      :record="record"
       class="annotation"
+      :dataset="dataset"
+      :record="record"
+      :visualTokens="visualTokens"
       :entities="getEntitiesByOrigin('annotation')"
     />
   </div>
@@ -49,6 +51,10 @@ export default {
       type: Object,
       required: true,
     },
+    visualTokens: {
+      type: Array,
+      required: true,
+    },
   },
   methods: {
     getEntitiesByOrigin(origin) {
diff --git a/frontend/components/token-classifier/results/TextSpan.vue b/frontend/components/token-classifier/results/TextSpan.vue
index 6aebaa3b65..f9a96ba381 100755
--- a/frontend/components/token-classifier/results/TextSpan.vue
+++ b/frontend/components/token-classifier/results/TextSpan.vue
@@ -22,6 +22,7 @@
       :class="['color_' + tag_color, { zindex3: showEntitiesSelector }]"
       :span="token"
       :dataset="dataset"
+      :record="record"
       @openTagSelector="openTagSelector"
       @removeEntity="removeEntity"
     /><span
@@ -31,11 +32,7 @@
       v-else
       v-for="(t, i) in token.tokens"
       :key="i"
-      v-html="
-        `${$highlightSearch(dataset.query.text, t.text)}${
-          token.hasSpaceAfter ? ' ' : ''
-        }`
-      "
+      v-html="visualizeToken(t)"
     ></span
     ><lazy-entities-selector
       :dataset="dataset"
@@ -155,6 +152,12 @@ export default {
       this.$emit("changeEntityLabel", token, entityLabel);
       this.showEntitiesSelector = false;
     },
+    visualizeToken(token) {
+      let text = token.highlighted
+        ? this.$htmlHighlightText(token.text)
+        : this.$htmlText(token.text);
+      return `${text}${token.hasSpaceAfter ? " " : ""}`;
+    },
   },
 };
 </script>
diff --git a/frontend/components/token-classifier/results/TextSpanStatic.vue b/frontend/components/token-classifier/results/TextSpanStatic.vue
index dec0f1aea5..abf7a5be72 100755
--- a/frontend/components/token-classifier/results/TextSpanStatic.vue
+++ b/frontend/components/token-classifier/results/TextSpanStatic.vue
@@ -22,6 +22,7 @@
       :class="['color_' + tag_color]"
       :span="token"
       :dataset="dataset"
+      :record="record"
     /><template v-else v-for="t in token.tokens">{{ t.text }}</template
     ><template>{{ token.hasSpaceAfter ? " " : "" }}</template>
   </span>
diff --git a/frontend/components/token-classifier/results/TextSpans.vue b/frontend/components/token-classifier/results/TextSpans.vue
index 1604014a6c..38966bc717 100755
--- a/frontend/components/token-classifier/results/TextSpans.vue
+++ b/frontend/components/token-classifier/results/TextSpans.vue
@@ -51,7 +51,6 @@
 
 <script>
 import { mapActions } from "vuex";
-import { indexOf, length } from "stringz";
 
 export default {
   props: {
@@ -66,6 +65,11 @@ export default {
       type: Object,
       required: true,
     },
+
+    visualTokens: {
+      type: Array,
+      required: true,
+    },
     origin: {
       type: String,
       required: true,
@@ -80,30 +84,6 @@ export default {
     };
   },
   computed: {
-    visualTokens() {
-      const recordHasEmoji = this.record.text.containsEmoji;
-      const { visualTokens } = this.record.tokens.reduce(
-        ({ visualTokens, startPosition }, token) => {
-          const start = recordHasEmoji
-            ? indexOf(this.record.text, token, startPosition)
-            : this.record.text.indexOf(token, startPosition);
-          const end = start + (recordHasEmoji ? length(token) : token.length);
-          const hasSpaceAfter = this.record.text.slice(end, end + 1) === " ";
-          return {
-            visualTokens: [
-              ...visualTokens,
-              { start, end, text: token, hasSpaceAfter: hasSpaceAfter },
-            ],
-            startPosition: end,
-          };
-        },
-        {
-          visualTokens: [],
-          startPosition: 0,
-        }
-      );
-      return Object.freeze(visualTokens);
-    },
     textSpans() {
       // TODO Simplify !!!
       const normalizedEntities = (entities, tokens) => {
diff --git a/frontend/components/token-classifier/results/TextSpansStatic.vue b/frontend/components/token-classifier/results/TextSpansStatic.vue
index c5b868cea1..9ec645855c 100755
--- a/frontend/components/token-classifier/results/TextSpansStatic.vue
+++ b/frontend/components/token-classifier/results/TextSpansStatic.vue
@@ -29,8 +29,6 @@
 </template>
 
 <script>
-import { indexOf, length } from "stringz";
-
 export default {
   props: {
     entities: {
@@ -44,36 +42,16 @@ export default {
       type: Object,
       required: true,
     },
+    visualTokens: {
+      type: Array,
+      required: true,
+    },
     origin: {
       type: String,
       required: true,
     },
   },
   computed: {
-    visualTokens() {
-      const recordHasEmoji = this.record.text.containsEmoji;
-      const { visualTokens } = this.record.tokens.reduce(
-        ({ visualTokens, startPosition }, token) => {
-          const start = recordHasEmoji
-            ? indexOf(this.record.text, token, startPosition)
-            : this.record.text.indexOf(token, startPosition);
-          const end = start + (recordHasEmoji ? length(token) : token.length);
-          const hasSpaceAfter = this.record.text.slice(end, end + 1) === " ";
-          return {
-            visualTokens: [
-              ...visualTokens,
-              { start, end, text: token, hasSpaceAfter: hasSpaceAfter },
-            ],
-            startPosition: end,
-          };
-        },
-        {
-          visualTokens: [],
-          startPosition: 0,
-        }
-      );
-      return visualTokens;
-    },
     textSpans() {
       // TODO Simplify !!!
       const normalizedEntities = (entities, tokens) => {
@@ -134,6 +112,7 @@ export default {
   },
 };
 </script>
+
 <style lang="scss" scoped>
 .content {
   &__input {
diff --git a/frontend/models/Common.js b/frontend/models/Common.js
index 05038c582f..b662986217 100644
--- a/frontend/models/Common.js
+++ b/frontend/models/Common.js
@@ -34,6 +34,7 @@ class BaseRecord {
     status,
     selected,
     event_timestamp,
+    search_keywords,
   }) {
     this.id = id;
     this.metadata = metadata;
@@ -43,6 +44,7 @@ class BaseRecord {
     this.status = status;
     this.selected = selected || false;
     this.event_timestamp = event_timestamp;
+    this.search_keywords = search_keywords || [];
   }
 
   recordTitle() {
diff --git a/frontend/plugins/highlight-search.js b/frontend/plugins/highlight-search.js
index 2087357e3d..4613bae63e 100644
--- a/frontend/plugins/highlight-search.js
+++ b/frontend/plugins/highlight-search.js
@@ -16,24 +16,65 @@
  */
 
 export default (context, inject) => {
-  const highlightSearch = function (query, text) {
-    const escapedText = text
+  const htmlText = function (text) {
+    return text
       .toString()
       .replace(/&/g, "&amp;")
       .replace(/</g, "&lt;")
       .replace(/>/g, "&gt;")
       .replace(/"/g, "&quot;")
       .replace(/'/g, "&#039;");
+  };
+
+  function htmlHighlightText(text) {
+    return `<span class="highlight-text">${htmlText(text)}</span>`;
+  }
+
+  const regexFromTerm = function (term) {
+    let q = term.replace(/[-[\]{}()*+?.,\\/^$|#\s]/g, "");
+    return new RegExp(q, "gi");
+  };
+
+  const highlightSearch = function (query, text) {
+    const escapedText = htmlText(text);
     if (!query) {
-      return escapedText;
+      return text;
     }
-    let q = query.replace(/[-[\]{}()*+?.,\\/^$|#\s]/g, "");
-    return escapedText
-      .toString()
-      .replace(
-        new RegExp(q, "gi"),
-        (match) => `<span class="highlight-text">${match}</span>`
-      );
+
+    return escapedText.replace(
+      regexFromTerm(query),
+      (match) => `<span class="highlight-text">${match}</span>`
+    );
+  };
+
+  const highlightKeywords = function (text, keywords) {
+    const sortedKeywords = ([...keywords] || []).sort(
+      (a, b) => b.length - a.length
+    );
+    text = htmlText(text);
+    sortedKeywords.forEach((keyword) => {
+      const regex = new RegExp(`\\b${keyword}\\b`, "gmi");
+      text = text.replace(regex, (match) => htmlHighlightText(match));
+    });
+
+    return text;
   };
+
+  const keywordsSpans = function (text, keywords) {
+    return (keywords || []).flatMap((keyword) => {
+      const regex = new RegExp(`\\b${keyword}\\b`, "gmi");
+      return [...text.matchAll(regex)].map((match) => {
+        return {
+          start: match.index,
+          end: match.index + match[0].length,
+        };
+      });
+    });
+  };
+
   inject("highlightSearch", highlightSearch);
+  inject("highlightKeywords", highlightKeywords);
+  inject("keywordsSpans", keywordsSpans);
+  inject("htmlText", htmlText);
+  inject("htmlHighlightText", htmlHighlightText);
 };
diff --git a/frontend/specs/components/text-classification/results/__snapshots__/RecordTextClassification.spec.js.snap b/frontend/specs/components/text-classification/results/__snapshots__/RecordTextClassification.spec.js.snap
index d4ef8bb6b7..fc02aedcc2 100644
--- a/frontend/specs/components/text-classification/results/__snapshots__/RecordTextClassification.spec.js.snap
+++ b/frontend/specs/components/text-classification/results/__snapshots__/RecordTextClassification.spec.js.snap
@@ -3,7 +3,7 @@
 exports[`RecordTextClassification renders properly 1`] = `
 <div class="record">
   <div class="record--left">
-    <recordinputs data="[object Object]" query-text="mock test"></recordinputs>
+    <recordinputs record="[object Object]"></recordinputs>
     <classifierexplorationarea dataset="[object Object]" record="[object Object]"></classifierexplorationarea>
     <!---->
   </div>
diff --git a/src/rubrix/server/tasks/commons/dao/dao.py b/src/rubrix/server/tasks/commons/dao/dao.py
index 97759baeb3..908e4a79ae 100644
--- a/src/rubrix/server/tasks/commons/dao/dao.py
+++ b/src/rubrix/server/tasks/commons/dao/dao.py
@@ -136,6 +136,10 @@ class DatasetRecordsDAO:
         rf"{__HIGHLIGHT_PRE_TAG__}(.+?){__HIGHLIGHT_POST_TAG__}"
     )
 
+    __HIGHLIGHT_PHRASE_PRE_PARSER_REGEX__ = re.compile(
+        rf"{__HIGHLIGHT_POST_TAG__}\s+{__HIGHLIGHT_PRE_TAG__}"
+    )
+
     @classmethod
     def get_instance(
         cls,
@@ -257,7 +261,7 @@ def search_records(
             "query": search.query or {"match_all": {}},
             "sort": sort_config,
             "aggs": aggregation_requests,
-            "highlight": self.__configure_query_highlight__(),
+            "highlight": self.__configure_query_highlight__(task=dataset.task),
         }
 
         try:
@@ -348,7 +352,7 @@ def scan_dataset(
         search = search or RecordSearch()
         es_query = {
             "query": search.query,
-            "highlight": self.__configure_query_highlight__(),
+            "highlight": self.__configure_query_highlight__(task=dataset.task),
         }
         docs = self._es.list_documents(
             dataset_records_index(dataset.id), query=es_query
@@ -356,15 +360,27 @@ def scan_dataset(
         for doc in docs:
             yield self.__esdoc2record__(doc)
 
-    def __esdoc2record__(self, doc: Dict[str, Any]):
+    def __esdoc2record__(
+        self,
+        doc: Dict[str, Any],
+        query: Optional[str] = None,
+        is_phrase_query: bool = True,
+    ):
         return {
             **doc["_source"],
             "id": doc["_id"],
-            "search_keywords": self.__parse_highlight_results__(doc),
+            "search_keywords": self.__parse_highlight_results__(
+                doc, query=query, is_phrase_query=is_phrase_query
+            ),
         }
 
     @classmethod
-    def __parse_highlight_results__(cls, doc: Dict[str, Any]) -> Optional[List[str]]:
+    def __parse_highlight_results__(
+        cls,
+        doc: Dict[str, Any],
+        query: Optional[str] = None,
+        is_phrase_query: bool = False,
+    ) -> Optional[List[str]]:
         highlight_info = doc.get("highlight")
         if not highlight_info:
             return None
@@ -373,8 +389,11 @@ def __parse_highlight_results__(cls, doc: Dict[str, Any]) -> Optional[List[str]]
         for content in highlight_info.values():
             if not isinstance(content, list):
                 content = [content]
-            for text in content:
-                search_keywords.extend(re.findall(cls.__HIGHLIGHT_VALUES_REGEX__, text))
+            text = " ".join(content)
+
+            if is_phrase_query:
+                text = re.sub(cls.__HIGHLIGHT_PHRASE_PRE_PARSER_REGEX__, " ", text)
+            search_keywords.extend(re.findall(cls.__HIGHLIGHT_VALUES_REGEX__, text))
         return list(set(search_keywords))
 
     def _configure_metadata_fields(self, index: str, metadata_values: Dict[str, Any]):
@@ -455,12 +474,19 @@ def get_dataset_schema(self, dataset: BaseDatasetDB) -> Dict[str, Any]:
         return self._es.__client__.indices.get_mapping(index=index_name)
 
     @classmethod
-    def __configure_query_highlight__(cls):
+    def __configure_query_highlight__(cls, task: TaskType):
+
         return {
             "pre_tags": [cls.__HIGHLIGHT_PRE_TAG__],
             "post_tags": [cls.__HIGHLIGHT_POST_TAG__],
             "require_field_match": False,
-            "fields": {"text": {}},
+            "fields": {
+                "text": {},
+                # TODO: `words` will be removed once the migration will be completed.
+                #  This configuration is included just for old datasets records
+                "words": {},
+                **({"inputs.*": {}} if task == TaskType.text_classification else {}),
+            },
         }
 
 
diff --git a/tests/functional_tests/test_log_for_text_classification.py b/tests/functional_tests/test_log_for_text_classification.py
index cf9bdb6d53..b34b41ef8a 100644
--- a/tests/functional_tests/test_log_for_text_classification.py
+++ b/tests/functional_tests/test_log_for_text_classification.py
@@ -72,7 +72,7 @@ def test_search_keywords(mocked_client):
             for keyword in keywords
         ]
     )
-    assert {"limit", "limits", "limited"} == top_keywords, top_keywords
+    assert {"limit", "limits", "limit?"} == top_keywords, top_keywords
 
 
 def test_log_records_with_empty_metadata_list(mocked_client):