You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

38 lines
1.2 KiB

3 months ago
"""
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.
"""