fix(llms) :Google Palm top_k value fix (#609)

pandasai/llm/base.py

@@ -375,7 +375,7 @@ class BaseGoogle(LLM):
 
     temperature: Optional[float] = 0
     top_p: Optional[float] = 0.8
-    top_k: Optional[float] = 0.3
+    top_k: Optional[int] = 40
     max_output_tokens: Optional[int] = 1000
 
     def _valid_params(self):
@@ -409,8 +409,8 @@ def _validate(self):
         if self.top_p is not None and not 0 <= self.top_p <= 1:
             raise ValueError("top_p must be in the range [0.0, 1.0]")
 
-        if self.top_k is not None and not 0 <= self.top_k <= 1:
-            raise ValueError("top_k must be in the range [0.0, 1.0]")
+        if self.top_k is not None and not 0 <= self.top_k <= 100:
+            raise ValueError("top_k must be in the range [0.0, 100.0]")
 
         if self.max_output_tokens is not None and self.max_output_tokens <= 0:
             raise ValueError("max_output_tokens must be greater than zero")

test_file.txt

@@ -0,0 +1 @@
+Hello, world!

tests/llms/test_google_palm.py

@@ -37,14 +37,14 @@ def test_params_setting(self):
             model="models/text-bison-001",
             temperature=0.5,
             top_p=1.0,
-            top_k=0.5,
+            top_k=50,
             max_output_tokens=64,
         )
 
         assert llm.model == "models/text-bison-001"
         assert llm.temperature == 0.5
         assert llm.top_p == 1.0
-        assert llm.top_k == 0.5
+        assert llm.top_k == 50
         assert llm.max_output_tokens == 64
 
     def test_validations(self, prompt):
@@ -69,14 +69,14 @@ def test_validations(self, prompt):
             GooglePalm(api_key="test", top_p=1.1).call(prompt, "World")
 
         with pytest.raises(
-            ValueError, match=re.escape("top_k must be in the range [0.0, 1.0]")
+            ValueError, match=re.escape("top_k must be in the range [0.0, 100.0]")
         ):
-            GooglePalm(api_key="test", top_k=-1).call(prompt, "World")
+            GooglePalm(api_key="test", top_k=-100).call(prompt, "World")
 
         with pytest.raises(
-            ValueError, match=re.escape("top_k must be in the range [0.0, 1.0]")
+            ValueError, match=re.escape("top_k must be in the range [0.0, 100.0]")
         ):
-            GooglePalm(api_key="test", top_k=1.1).call(prompt, "World")
+            GooglePalm(api_key="test", top_k=110).call(prompt, "World")
 
         with pytest.raises(
             ValueError, match=re.escape("max_output_tokens must be greater than zero")