choices

evaluations/models.py

@@ -29,14 +29,8 @@ class RemoteLLM():
     self.client = InferenceClient(api_key=token)
 
   def predict(self, response_type, prompt):
-    for message in self.client.chat_completion(
-        model=self.model_name,
+    message = self.client.chat_completion(
+        model=self.model_name, max_tokens=500, stream=False,
         messages=[{"role": "system", "content": system_messages[response_type]}, 
-                  {"role": "user", "content": prompt}],
-        max_tokens=500,
-        stream=False,
-        ):
-        return message['choices']
-
-    return ""
-
+                  {"role": "user", "content": prompt}])
+    return message['choices'][0]