zsk / model_utils.py
ikun520's picture
Update model_utils.py
a3dd631 verified
from langchain_community.embeddings import HuggingFaceEmbeddings
from sparkai.llm.llm import ChatSparkLLM
from sparkai.core.messages import ChatMessage
# 星火 API 配置信息
SPARKAI_URL = 'wss://spark-api.xf-yun.com/v1.1/chat'
SPARKAI_APP_ID = 'ebf052b9'
SPARKAI_API_SECRET = 'NzIwNzYyOWM0ZmM3OTk3YWVmMDlmZTVj'
SPARKAI_API_KEY = '930eb4722c8d5bc3254779893f7a34e4'
SPARKAI_DOMAIN = 'general'
def get_chat_model(model="spark-gpt", temperature=0):
"""
获取星火聊天模型
"""
return ChatSparkLLM(
spark_api_url=SPARKAI_URL,
spark_app_id=SPARKAI_APP_ID,
spark_api_key=SPARKAI_API_KEY,
spark_api_secret=SPARKAI_API_SECRET,
spark_llm_domain=SPARKAI_DOMAIN,
streaming=False
)
# 使用模型名称初始化 HuggingFaceEmbeddings
embeddings_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
def matter(txt):
"""
调用星火 API 生成结果
"""
spark = get_chat_model()
messages = [ChatMessage(
role="user",
content=txt
)]
a = spark.generate([messages])
return a.generations[0][0].text