35 lines
1.3 KiB
Python
35 lines
1.3 KiB
Python
|
from modelscope import AutoModelForCausalLM, AutoTokenizer
|
||
|
|
|
||
|
|
class QwenChatbot:
|
||
|
|
def __init__(self, model_name="Qwen/Qwen3-0.6B"):
|
||
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
|
||
|
|
self.model = AutoModelForCausalLM.from_pretrained(model_name)
|
||
|
|
self.history = []
|
||
|
|
|
||
|
|
def generate_response(self, user_input):
|
||
|
|
messages = self.history + [{"role": "system", "content": "你是一个翻译模型，请将输入翻译成中文。"},{"role": "user", "content": user_input}]
|
||
|
|
|
||
|
|
text = self.tokenizer.apply_chat_template(
|
||
|
|
messages,
|
||
|
|
tokenize=False,
|
||
|
|
add_generation_prompt=True,
|
||
|
|
enable_thinking=False # Setting enable_thinking=False disables thinking mode
|
||
|
|
|
||
|
|
)
|
||
|
|
|
||
|
|
inputs = self.tokenizer(text, return_tensors="pt")
|
||
|
|
response_ids = self.model.generate(**inputs, max_new_tokens=32768)[0][len(inputs.input_ids[0]):].tolist()
|
||
|
|
response = self.tokenizer.decode(response_ids, skip_special_tokens=True)
|
||
|
|
|
||
|
|
# Update history
|
||
|
|
self.history.append({"role": "user", "content": user_input})
|
||
|
|
self.history.append({"role": "assistant", "content": response})
|
||
|
|
|
||
|
|
return response
|
||
|
|
|
||
|
|
# Example Usage
|
||
|
|
if __name__ == "__main__":
|
||
|
|
chatbot = QwenChatbot()
|
||
|
|
|
||
|
|
print(chatbot.generate_response("Hello"))
|