From 8d8eae158c91f3b1ea2e50a3e4f309feb77ec8f2 Mon Sep 17 00:00:00 2001
From: Limbo <57550280+Alu-XXX@users.noreply.github.com>
Date: Mon, 13 May 2024 17:26:09 +0800
Subject: [PATCH] Update emotion.py

---
 src/blackbox/emotion.py | 23 +++++++++++++++++------
 1 file changed, 17 insertions(+), 6 deletions(-)

diff --git a/src/blackbox/emotion.py b/src/blackbox/emotion.py
index 80b6966..af0fce1 100644
--- a/src/blackbox/emotion.py
+++ b/src/blackbox/emotion.py
@@ -11,7 +11,7 @@ from injector import singleton
 @singleton
 class Emotion(Blackbox):
 
-    def __init__(self, model_name, model_url) -> None:
+    def __init__(self, model_name:str = "Mistral-7B-Instruct-v0.2", model_url:str = "http://120.196.116.194:48892") -> None:
         self.model = LMDeployClient(
             model_name=model_name,
             url=model_url,
@@ -29,10 +29,21 @@ class Emotion(Blackbox):
         data = args[0]
         return isinstance(data, str)
     
-    def processing(self, text: str) -> int:
-        text = "Please use one word to infer the emotion of the following passage:\n" + text + "\nJust print out that signle word pls."
-        text = [{'role': 'user', 'content': text}]
-        return self.model.stream_chat(text)
+    def processing(self, text: str) -> str:
+
+        response = ""
+
+        for text in self.model.stream_chat(text):
+            print("------------")
+            print(text[1])
+            print("------------")
+            # print(type(text))
+
+        response = text[1]
+        # print("************0")
+        # print(response)
+
+        return response
     
     async def fast_api_handler(self, request) -> Response:
         try:
@@ -45,5 +56,5 @@ class Emotion(Blackbox):
         text = "Please use one word to infer the emotion of the following passage:\n" + text + "\nJust print out that signle word pls."
         text = [{'role': 'user', 'content': text}]
         sentiment = self.processing(text)
-        return JSONResponse(content={"sentiment": sentiment }, status_code=status.HTTP_200_OK)
+        return JSONResponse(content={"sentiment": sentiment}, status_code=status.HTTP_200_OK)