47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
import logging
|
||
from typing import Tuple
|
||
import requests
|
||
from ollama import Client, ResponseError
|
||
import tiktoken
|
||
import random
|
||
from Utils.AIGCLog import AIGCLog
|
||
|
||
class AICore:
|
||
modelMaxTokens = 128000
|
||
# 初始化 DeepSeek 使用的 Tokenizer (cl100k_base)
|
||
encoder = tiktoken.get_encoding("cl100k_base")
|
||
logger = AIGCLog(name = "AIGC", log_file = "aigc.log")
|
||
|
||
|
||
def __init__(self, model):
|
||
#初始化ollama客户端
|
||
self.ollamaClient = Client(host='http://localhost:11434', headers={'x-some-header': 'some-value'})
|
||
self.modelName = model
|
||
response = self.ollamaClient.show(model)
|
||
modelMaxTokens = response.modelinfo['qwen2.context_length']
|
||
|
||
def getPromptToken(self, prompt)-> int:
|
||
tokens = self.encoder.encode(prompt)
|
||
return len(tokens)
|
||
|
||
def generateAI(self, promptStr: str) -> Tuple[bool, str]:
|
||
try:
|
||
response = self.ollamaClient.generate(
|
||
model = self.modelName,
|
||
stream = False,
|
||
prompt = promptStr,
|
||
options={
|
||
"temperature": random.uniform(1.0, 1.5),
|
||
"repeat_penalty": 1.2, # 抑制重复
|
||
"top_p": random.uniform(0.7, 0.95),
|
||
"num_ctx": 4096, # 上下文长度
|
||
}
|
||
)
|
||
return True, response.response
|
||
except ResponseError as e:
|
||
if e.status_code == 503:
|
||
print("🔄 服务不可用，5秒后重试...")
|
||
return False,"ollama 服务不可用"
|
||
except Exception as e:
|
||
print(f"🔥 未预料错误: {str(e)}")
|
||
return False, "未预料错误" |