38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
"""
|
|
This file is designed to invoke OpenAI GPT-4
|
|
to generate possible prompt words for specifying manim code.
|
|
"""
|
|
|
|
import os
|
|
import openai
|
|
from openai import OpenAI
|
|
from langchain_core.prompts import ChatPromptTemplate
|
|
from langchain_openai import ChatOpenAI
|
|
|
|
# Define the GPT-4 model endpoint
|
|
llm = ChatOpenAI(model="gpt-4", openai_api_key=os.environ["OPENAI_API_KEY"])
|
|
|
|
# input and output
|
|
chain = ChatPromptTemplate.from_messages(
|
|
[
|
|
(
|
|
"system",
|
|
"If you are a user who is going to use GPT to generate manim code, write the prompt words used by the role according to the code I gave you, write in the same paragraph, do not need to step by step analysis, do not explain your output in the first paragraph:",
|
|
),
|
|
(
|
|
"user",
|
|
"{input}",
|
|
),
|
|
]
|
|
)
|
|
|
|
reply = chain.format_messages(
|
|
input="from manim import * class MyScene(Scene): def construct(self): circle = Circle(radius=2, color=BLUE) self.add(circle)"
|
|
)
|
|
|
|
print(llm(reply).content)
|
|
"""
|
|
e.g output
|
|
Generate the Manim code for me that creates a scene with a blue circle of radius 2 centered at the origin. The circle should appear gradually on the screen from left to right. Include comments in the code to explain each step.
|
|
"""
|