fixing ml model feature tab (datahub-project#4769)

maggiehays · Aug 1, 2022 · e4389b7 · e4389b7
1 parent cdfe3d1
commit e4389b7
Showing 1 changed file with 2 additions and 5 deletions.
diff --git a/datahub-web-react/src/app/entity/mlModel/profile/MlModelFeaturesTab.tsx b/datahub-web-react/src/app/entity/mlModel/profile/MlModelFeaturesTab.tsx
@@ -1,18 +1,15 @@
 import React from 'react';
 
 import { EntityType } from '../../../../types.generated';
-import { useEntityRegistry } from '../../../useEntityRegistry';
 import { useBaseEntity } from '../../shared/EntityContext';
 import { GetMlModelQuery } from '../../../../graphql/mlModel.generated';
 import { EntityList } from '../../shared/tabs/Entity/components/EntityList';
 
 export default function MlModelFeaturesTab() {
     const entity = useBaseEntity() as GetMlModelQuery;
-    const entityRegistry = useEntityRegistry();
 
     const model = entity && entity.mlModel;
-    const featureTables = model?.features?.relationships.map((relationship) => relationship.entity);
+    const features = model?.features?.relationships.map((relationship) => relationship.entity);
 
-    const title = `Part of ${entityRegistry.getEntityName(EntityType.MlfeatureTable)}`;
-    return <EntityList title={title} type={EntityType.MlfeatureTable} entities={featureTables || []} />;
+    return <EntityList title="Features used" type={EntityType.Mlfeature} entities={features || []} />;
 }