54 lines
1.8 KiB
Python
54 lines
1.8 KiB
Python
"""
|
|
This script is designed for interacting with a local GLM3 AI model using the `ChatGLM3` class
|
|
from the `langchain_community` library. It facilitates continuous dialogue with the GLM3 model.
|
|
|
|
1. Start the Local Model Service: Before running this script, you need to execute the `api_server.py` script
|
|
to start the GLM3 model's service.
|
|
2. Run the Script: The script includes functionality for initializing the LLMChain object and obtaining AI responses,
|
|
allowing the user to input questions and receive AI answers.
|
|
"""
|
|
from langchain.chains import LLMChain
|
|
from langchain.prompts import PromptTemplate
|
|
from langchain.schema.messages import HumanMessage, SystemMessage, AIMessage
|
|
from langchain_community.llms.chatglm3 import ChatGLM3
|
|
|
|
|
|
def initialize_llm_chain(messages: list):
|
|
template = "{input}"
|
|
prompt = PromptTemplate.from_template(template)
|
|
|
|
endpoint_url = "http://127.0.0.1:8000/v1/chat/completions"
|
|
llm = ChatGLM3(
|
|
endpoint_url=endpoint_url,
|
|
max_tokens=8096,
|
|
prefix_messages=messages,
|
|
top_p=0.9,
|
|
)
|
|
return LLMChain(prompt=prompt, llm=llm)
|
|
|
|
|
|
def get_ai_response(llm_chain, user_message):
|
|
ai_response = llm_chain.invoke({"input": user_message})
|
|
return ai_response
|
|
|
|
|
|
def continuous_conversation():
|
|
messages = [
|
|
SystemMessage(content="You are an intelligent AI assistant, named ChatGLM3."),
|
|
]
|
|
while True:
|
|
user_input = input("Human (or 'exit' to quit): ")
|
|
if user_input.lower() == 'exit':
|
|
break
|
|
llm_chain = initialize_llm_chain(messages=messages)
|
|
ai_response = get_ai_response(llm_chain, user_input)
|
|
print("ChatGLM3: ", ai_response["text"])
|
|
messages += [
|
|
HumanMessage(content=user_input),
|
|
AIMessage(content=ai_response["text"]),
|
|
]
|
|
|
|
|
|
if __name__ == "__main__":
|
|
continuous_conversation()
|