AI_NEWS/api_ollama.py

# -*- coding: utf-8 -*-
import time
import subprocess
import sys
import timeit
from ollama import Client as oClient

'''
可选模型:
qwen:1.8b
qwen2:1.5b
qwen2.5:3b
gemma2:2b
'''
model_list = [
    'qwen:1.8b',
    'qwen2:1.5b',
    'qwen2.5:3b',
    'gemma2:2b'
]


class ChatBot:
    def __init__(self, host, messages, model='qwen:1.8b', temperature=0.4):
        self.client = oClient(host=host)
        self.model = model
        self.messages = messages
        self.temperature = temperature

    def start_chat(self):
        print(f'use model: {self.model}')
        try:
            response_iter = self.client.chat(model=self.model,
                                             messages=[
                                                 {'role': 'user', 'content': '你是一个新闻整理专员'},
                                                 {'role': 'user', 'content': self.messages}
                                             ],
                                             options={"temperature": self.temperature},
                                             stream=False)
            return response_iter['message']['content']
        except Exception as e:
            print(f"\n发生错误: {e}")


# if __name__ == "__main__":
#     for m in model_list:
#         C = ChatBot('http://erhe.top:27381', 'hello,你好呀', m)
#         start_time = time.time()
#
#         response_context = C.start_chat()
#         print(response_context)
#
#         end_time = time.time()
#         run_time = end_time - start_time
#         print(f"程序运行时间：{run_time} 秒\n")