wehub-resource-sync
24 lines
692 B
Python
24 lines
692 B
Python
import os
|
|
|
|
from langchain.llms import OpenAI
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
|
|
import mlflow
|
|
|
|
# Ensure the OpenAI API key is set in the environment
|
|
assert "OPENAI_API_KEY" in os.environ, "Please set the OPENAI_API_KEY environment variable."
|
|
|
|
# Initialize the OpenAI model and the prompt template
|
|
llm = OpenAI()
|
|
|
|
# Create the LLMChain with the specified model and prompt
|
|
chain = llm | StrOutputParser()
|
|
|
|
with mlflow.start_run() as run:
|
|
model_info = mlflow.langchain.log_model(chain, name="model")
|
|
|
|
loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
|
|
|
|
for chunk in loaded_model.predict_stream("Count to 10. E.g., 1, 2, 3, ..."):
|
|
print(chunk, end="|")
|