Fix lint error

PiperOrigin-RevId: 661241075 Change-Id: I6908050e002b0663f2abf181126f3e74f9bb957c
google-deepmind · Aug 9, 2024 · f86f4de · f86f4de
1 parent 0ab7ced
commit f86f4de
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/concordia/language_model/ollama_model.py b/concordia/language_model/ollama_model.py
@@ -91,7 +91,7 @@ def sample_text(
         options={'stop': terminators},
         keep_alive='10m',
     )
-    result = response['response']
+    result = response.response
 
     if self._measurements is not None:
       self._measurements.publish_datum(
@@ -126,9 +126,8 @@ def sample_choice(
           format='json',
           keep_alive='10m',
       )
-      json_data = response
       try:
-        json_data_response = json.loads(json_data['response'])
+        json_data_response = json.loads(response.response)
       except json.JSONDecodeError:
         continue
       sample_or_none = json_data_response.get('choice', None)