diff --git a/frontend/src/components/nlp/NlpValueDialog.tsx b/frontend/src/components/nlp/NlpValueDialog.tsx
index 9eb0341e..2524805e 100644
--- a/frontend/src/components/nlp/NlpValueDialog.tsx
+++ b/frontend/src/components/nlp/NlpValueDialog.tsx
@@ -20,10 +20,11 @@ import { ContentItem } from "@/app-components/dialogs/layouts/ContentItem";
 import { Input } from "@/app-components/inputs/Input";
 import MultipleInput from "@/app-components/inputs/MultipleInput";
 import { useCreate } from "@/hooks/crud/useCreate";
+import { useGet } from "@/hooks/crud/useGet";
 import { useUpdate } from "@/hooks/crud/useUpdate";
 import { DialogControlProps } from "@/hooks/useDialog";
 import { useToast } from "@/hooks/useToast";
-import { EntityType } from "@/services/types";
+import { EntityType, Format } from "@/services/types";
 import { INlpValue, INlpValueAttributes } from "@/types/nlp-value.types";
 
 export type TNlpValueAttributesWithRequiredExpressions = INlpValueAttributes & {
@@ -44,11 +45,16 @@ export const NlpValueDialog: FC<NlpValueDialogProps> = ({
   const { t } = useTranslation();
   const { toast } = useToast();
   const { query } = useRouter();
+  const { refetch: refetchEntity } = useGet(data?.entity || String(query.id), {
+    entity: EntityType.NLP_ENTITY,
+    format: Format.FULL,
+  });
   const { mutateAsync: createNlpValue } = useCreate(EntityType.NLP_VALUE, {
     onError: () => {
       toast.error(t("message.internal_server_error"));
     },
     onSuccess(data) {
+      refetchEntity();
       closeDialog();
       toast.success(t("message.success_save"));
       callback?.(data);
diff --git a/frontend/src/components/nlp/components/NlpTrainForm.tsx b/frontend/src/components/nlp/components/NlpTrainForm.tsx
index 81cd097f..0d6bd88b 100644
--- a/frontend/src/components/nlp/components/NlpTrainForm.tsx
+++ b/frontend/src/components/nlp/components/NlpTrainForm.tsx
@@ -23,7 +23,7 @@ import {
   RadioGroup,
   Typography,
 } from "@mui/material";
-import { FC, useCallback, useMemo, useState } from "react";
+import { FC, useCallback, useEffect, useMemo, useState } from "react";
 import { Controller, useFieldArray, useForm } from "react-hook-form";
 import { useTranslation } from "react-i18next";
 import { useQuery } from "react-query";
@@ -65,71 +65,40 @@ const NlpDatasetSample: FC<NlpDatasetSampleProps> = ({
     {
       hasCount: false,
     },
-    {
-      onSuccess(entities) {
-        // By default append trait entities
-        if (!sample) {
-          removeTraitEntity();
-          (entities || [])
-            .filter(({ lookups }) => lookups.includes("trait"))
-            .forEach(({ name }) => {
-              appendTraitEntity({
-                entity: name,
-                value: "",
-              });
-            });
-        }
-      },
-    },
   );
   const getNlpValueFromCache = useGetFromCache(EntityType.NLP_VALUE);
-  // Default trait entities to append to the form
-  const defaultTraitEntities = useMemo(() => {
-    if (!sample || !entities) return [];
-
-    const traitEntities = entities.filter(({ lookups }) =>
-      lookups.includes("trait"),
-    );
-    const sampleTraitEntities = sample.entities.filter(
-      (e) => "start" in e && typeof e.start === "undefined",
-    );
-
-    if (sampleTraitEntities.length === traitEntities.length) {
-      return sampleTraitEntities;
-    }
-
-    const sampleEntityNames = new Set(sampleTraitEntities.map((e) => e.entity));
-    const missingEntities = traitEntities
-      .filter(({ name }) => !sampleEntityNames.has(name))
-      .map(({ name }) => ({
-        entity: name,
-        value: "",
-      }));
-
-    return [...sampleTraitEntities, ...missingEntities];
-  }, [entities, sample]);
+  // eslint-disable-next-line react-hooks/exhaustive-deps
+  const defaultValues: INlpSampleFormAttributes = useMemo(
+    () => ({
+      type: sample?.type || NlpSampleType.train,
+      text: sample?.text || "",
+      language: sample?.language || null,
+      traitEntities: (entities || [])
+        .filter(({ lookups }) => {
+          return lookups.includes("trait");
+        })
+        .map((e) => {
+          return {
+            entity: e.name,
+            value: sample
+              ? sample.entities.find(({ entity }) => entity === e.name)?.value
+              : "",
+          } as INlpDatasetTraitEntity;
+        }),
+      keywordEntities: (sample?.entities || []).filter(
+        (e) => "start" in e && typeof e.start === "number",
+      ) as INlpDatasetKeywordEntity[],
+    }),
+    [sample, entities],
+  );
   const { handleSubmit, control, register, reset, setValue, watch } =
     useForm<INlpSampleFormAttributes>({
-      defaultValues: {
-        type: sample?.type || NlpSampleType.train,
-        text: sample?.text || "",
-        language: sample?.language,
-        traitEntities: defaultTraitEntities,
-        keywordEntities:
-          sample?.entities.filter(
-            (e) => "start" in e && typeof e.start === "number",
-          ) || [],
-      },
+      defaultValues,
     });
   const currentText = watch("text");
   const currentType = watch("type");
   const { apiClient } = useApiClient();
-  const {
-    fields: traitEntities,
-    append: appendTraitEntity,
-    update: updateTraitEntity,
-    remove: removeTraitEntity,
-  } = useFieldArray({
+  const { fields: traitEntities, update: updateTraitEntity } = useFieldArray({
     control,
     name: "traitEntities",
   });
@@ -187,16 +156,14 @@ const NlpDatasetSample: FC<NlpDatasetSampleProps> = ({
   } | null>(null);
   const onSubmitForm = (form: INlpSampleFormAttributes) => {
     submitForm(form);
-    reset({
-      type: form?.type || NlpSampleType.train,
-      text: "",
-      language: form?.language,
-      traitEntities: defaultTraitEntities,
-      keywordEntities: [],
-    });
     refetchEntities();
   };
 
+  useEffect(() => {
+    reset(defaultValues);
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [JSON.stringify(defaultValues)]);
+
   return (
     <Box className="nlp-train" sx={{ position: "relative", p: 2 }}>
       <form onSubmit={handleSubmit(onSubmitForm)}>
diff --git a/frontend/src/components/nlp/components/NlpValues.tsx b/frontend/src/components/nlp/components/NlpValues.tsx
index 8ba2a9f9..a7c2415e 100644
--- a/frontend/src/components/nlp/components/NlpValues.tsx
+++ b/frontend/src/components/nlp/components/NlpValues.tsx
@@ -142,8 +142,6 @@ export const NlpValues = ({ entityId }: { entityId: string }) => {
 
   return (
     <Grid container gap={2} flexDirection="column">
-      {/* <PageHeader title={t("title.nlp_train")} icon={faGraduationCap} /> */}
-
       <Slide direction={direction} in={true} mountOnEnter unmountOnExit>
         <Grid item xs={12}>
           <Box sx={{ padding: 1 }}>