diff --git a/cldk/models/treesitter/models.py b/cldk/models/treesitter/models.py
index 2ef865a..349876d 100644
--- a/cldk/models/treesitter/models.py
+++ b/cldk/models/treesitter/models.py
@@ -1,6 +1,6 @@
 from dataclasses import dataclass
-from typing import List, Tuple
-
+from typing import Dict, List, Tuple
+from ipdb import set_trace
 from tree_sitter import Node
 
 
@@ -27,8 +27,9 @@ class Capture:
         node: Node
         name: str
 
-    def __init__(self, captures: List[Tuple[Node, str]]):
-        self.captures = [self.Capture(node=node, name=text) for node, text in captures]
+    def __init__(self, captures: Dict[str, List[Node]]):
+        for capture_name, captures in captures.items():
+            self.captures = [self.Capture(node=node, name=capture_name) for node in captures]
 
     def __getitem__(self, index: int) -> Capture:
         """Get the capture at the specified index.
diff --git a/pyproject.toml b/pyproject.toml
index 8077024..49b8b1c 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cldk"
-version = "0.1.0"
+version = "0.1.1"
 description = "codellm-devkit: A python library for seamless integration with LLMs."
 authors = ["Rahul Krishna <i.m.ralk@gmail.com>", "Rangeet Pan <rangeet.pan@ibm.com>", "Saurabh Sinhas <sinhas@us.ibm.com>",
             "Raju Pavuluri <pavuluri@us.ibm.com>"]