diff --git a/lua/base46/integrations/treesitter.lua b/lua/base46/integrations/treesitter.lua
index 1e34599..f90b4b1 100644
--- a/lua/base46/integrations/treesitter.lua
+++ b/lua/base46/integrations/treesitter.lua
@@ -1,7 +1,8 @@
 local theme = require("base46").get_theme_tb "base_16"
 
 return {
-  TSAnnotation = {
+  -- `@annotation` is not one of the default capture group, should we keep it 
+  ["@annotation"] = { 
     fg = theme.base0F,
   },
 
@@ -112,7 +113,7 @@ return {
   -- @tag
 
   ["@tag.attribute"] = {
-    link = "TSProperty",
+    link = "@property",
   },
 
   ["@tag.delimiter"] = {
diff --git a/lua/base46/themes/chocolate.lua b/lua/base46/themes/chocolate.lua
index fdf327a..1d4cacc 100644
--- a/lua/base46/themes/chocolate.lua
+++ b/lua/base46/themes/chocolate.lua
@@ -63,7 +63,7 @@ M.polish_hl = {
   Operator = { fg = M.base_30.blue },
   ["@attribute"] = { fg = M.base_30.cyan },
   ["@punctuation.bracket"] = { fg = M.base_16.base06 },
-  ["@parenthesis"] = { link = "TSPunctBracket" },
+  ["@parenthesis"] = { link = "@punctuation.bracket" },
   ["@parameter"] = { fg = M.base_30.green },
   ["@function.builtin"] = { fg = M.base_30.yellow },
 }