pease

generators/model.py

@@ -123,6 +123,7 @@ class Samba():
     
     def generate_chat(self, messages: List[Message], max_tokens: int = 1024, temperature: float = 0.2, num_comps: int = 1) -> Union[List[str], str]: 
         resps = []
+        
         for i in range(num_comps):
             payload = {
                 "inputs": [dataclasses.asdict(message) for message in messages],
@@ -145,19 +146,17 @@ class Samba():
                 "Content-Type": "application/json"
             }
             post_response = requests.post(url, json=payload, headers=headers, stream=True)
-
             response_text = ""
             for line in post_response.iter_lines():
                 if line.startswith(b"data: "):
                     data_str = line.decode('utf-8')[6:]
                     try:
                         line_json = json.loads(data_str)
-                        content = line_json.get("stream_token", "")
+                        content = line_json['0'].get("stream_token", "")
                         if content:
                             response_text += content
                     except json.JSONDecodeError as e:
                         pass
-            resps.append(response_text)
             
         if num_comps == 1:
             return resps[0]