Update networks/huggingface_streamer.py

networks/huggingface_streamer.py

@@ -176,3 +176,34 @@ class HuggingfaceStreamer:
 
         if not is_finished:
             yield self.message_outputer.output(content="", content_type="Finished")
+
+
+    def ollama_return_generator(self, stream_response):
+        is_finished = False
+        line_count = 0
+        for line in stream_response.iter_lines():
+            if line:
+                line_count += 1
+            else:
+                continue
+
+            content = self.parse_line(line)
+
+            if content.strip().endswith(self.stop_sequences):
+                content_type = "Finished"
+                logger.success("\n[Finished]")
+                is_finished = True
+            else:
+                content_type = "Completions"
+                if line_count == 1:
+                    content = content.lstrip()
+
+            logger.back(content, end="")
+
+            output = self.message_outputer.output(
+                content=content, content_type=content_type
+            )
+            yield "{" + output + "}"
+
+        if not is_finished:
+            yield "{" + self.message_outputer.output(content="", content_type="Finished") + "}"