From f86f4de108a9aa2d3aa8698005dea234685c3b7c Mon Sep 17 00:00:00 2001
From: John Agapiou <jagapiou@google.com>
Date: Fri, 9 Aug 2024 06:29:51 -0700
Subject: [PATCH] Fix lint error

PiperOrigin-RevId: 661241075
Change-Id: I6908050e002b0663f2abf181126f3e74f9bb957c
---
 concordia/language_model/ollama_model.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/concordia/language_model/ollama_model.py b/concordia/language_model/ollama_model.py
index 1157f3a3..7fa439b2 100644
--- 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)