Gemma[Chat]LocalKaggle class always loads the default model

[aczyzewski]

Hi everyone,

Version:

langchain-google-vertexai        1.0.5

I'm having issues with the GemmaLocalKaggle and GemmaChatLocalKaggle classes. They always use the model named "gemma_2b_en", even if I try to use a different one (e.g. gemma_1.1_instruct_2b_en).

Just for testing purposes I hard-coded the model name inside the library's code and everything worked fine. Here's the line that loads the default value instead of the model name that is passed as an arg.

This code shows the problem:

from langchain_google_vertexai import GemmaChatLocalKaggle

keras_backend: str = "jax"
model_name = "gemma_1.1_instruct_2b_en"
llm = GemmaChatLocalKaggle(model_name=model_name, keras_backend=keras_backend)
output = llm.invoke("What is the meaning of life?", max_tokens=30)
print(output)

And here are the logs that I'm getting.

Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'task.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'config.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'model.weights.h5' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'metadata.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'preprocessor.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...
Attaching 'tokenizer.json' from model 'keras/gemma/keras/gemma_2b_en/2' to your Colab notebook...

So as you can see, it loaded the wrong model. Is there an easy fix for that?

Update

It appears that the argument here is only set correctly when the model keyword argument is passed. However, all inheriting classes use model_name. This is a bit inconsistent with the Gemma documentation.

Workaround: you can initialise the class with two arguments: model and model_name, both set to the same value.

End update

Thanks, Adam

[lkuligin]

Would you like to send a PR fixing the issue, please (since you've already found the problem!)?

[aczyzewski]

Sure! I will submit a PR in the meantime.