Upload folder using huggingface_hub

ChatWorld/ChatWorld.py

@@ -66,8 +66,9 @@ class ChatWorld:
         self.model_role_nickname = role_nick_name
 
     def getSystemPrompt(self, role_name, role_nick_name):
-        assert self.model_role_name and self.model_role_nickname, "Please set model role name first"
-        return self.prompt.render(model_role_name=self.model_role_name, model_role_nickname=self.model_role_nickname, role_name=role_name, role_nickname=role_nick_name)
+        assert self.model_role_name, "Please set model role name first"
+
+        return {"role": "system", "content": self.prompt.render(model_role_name=self.model_role_name, model_role_nickname=self.model_role_nickname, role_name=role_name, role_nickname=role_nick_name)}
 
     def chat(self, user_role_name: str, text: str, user_role_nick_name: str = None, use_local_model=False):
         message = [self.getSystemPrompt(

ChatWorld/models.py

@@ -3,9 +3,23 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 
 class qwen_model:
     def __init__(self, model_name):
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-        self.model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto", trust_remote_code=True).eval()
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name, trust_remote_code=True)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, device_map="auto", trust_remote_code=True).eval()
 
     def get_response(self, message):
-        self.tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True)
-        return "test"
+        message = self.tokenizer.apply_chat_template(
+            message, tokenize=False, add_generation_prompt=True)
+        model_inputs = self.tokenizer([message], return_tensors="pt")
+        generated_ids = self.model.generate(
+            model_inputs.input_ids,
+            max_new_tokens=512
+        )
+        generated_ids = [
+            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
+        ]
+
+        response = self.tokenizer.batch_decode(
+            generated_ids, skip_special_tokens=True)[0]
+        return response