BaseEmbeddings.embed should be treated as callable

[ascendant512]

I tried to write a BaseEmbeddings implementation and it crashed due to this line:

chonkie/src/chonkie/chunker/base.py

Line 34 in 9904935

elif inspect.isfunction(tokenizer_or_token_counter):

It should be elif callable(tokenizer_or_token_counter):
The test fails because BaseEmbeddings.embed is a method, which inspect.isfunction does not consider to be a function.

The symptom of the crash was in the line Tokenizer backend {str(type(self.tokenizer))} not supported where it referred to my class and its method embed.

FYI, here is my class, I don't really know yet if it's correct, but I know that my change allowed it to run:

@dataclass
class OllamaEmbeddings(BaseEmbeddings):
	client: PatchedClient
	model: str

	def count_tokens(self, text: str) -> int:
		tokenized = self.client.tokenize(model=self.model, prompt=text)
		return len(tokenized.tokens)

	def embed(self, text: str) -> ndarray:
		embedded = self.client.embed(self.model, text)
		return array(embedded.embeddings)

	@property
	def dimension(self) -> int:
		return 2

	def embed_batch(self, texts: List[str]) -> List[ndarray]:
		return super().embed_batch(texts)

	def count_tokens_batch(self, texts: List[str]) -> List[int]:
		return super().count_tokens_batch(texts)

	def similarity(self, u: ndarray, v: ndarray) -> float:
		numerator = matmul(u, v.T) #i do not understand why it needs .T
		denominator = norm(u) * norm(v)
		return numerator / denominator

[bhavnicksm]

Hey @ascendant512!

Thanks for opening an issue! 😊

Chonkie had that line as a callable before but the reason why it was changed to the current method is because all tokenizers are technically callable since they implement the __call__ fn in their classes, which causes that conditional to error out.

If you can do one check, that would be great! Could you see by adding the get_tokenizer_or_token_counter() funtion into your class as:

def get_tokenizer_or_token_counter(self): 
    return self.count_tokens

Since this returns a function that counts tokens, it should by-pass the issue.

Please let me know the results!

Thanks! 😊

[ascendant512]

Hi, thanks for the quick reply!

I haven't use the other tokenizers yet, so I didn't know how they worked. However, I can confirm that your suggestion also didn't work. I get a different error:
ValueError: Tokenizer backend <class 'method'> not supported
I think the underlying issue is the same, relating to inspect.isfunction only returning true for a function that's not bound to a class. This does work as a substitute for callable, with or without with your suggested edit:
inspect.ismethod(tokenizer_or_token_counter)

[bhavnicksm]

Hey @ascendant512!

Sorry for the delay and thanks for trying it out; ismethod also sounds like a good option. I'll try evaluating the options we have and get a fix pushed ASAP.

Thanks~ 😊