|
import os |
|
import streamlit as st |
|
from keyfile import secret_value |
|
from langchain_openai import ChatOpenAI |
|
from langchain.prompts import PromptTemplate |
|
from langchain.chains import LLMChain, SequentialChain |
|
|
|
os.environ["OPENAI_API_KEY"] = secret_value |
|
|
|
def main(): |
|
st.title("Chain Based Story Generator") |
|
|
|
input1 = st.text_input("An inspirational object to start a story") |
|
input2 = st.text_input("A tone or a writer's style to shape the story") |
|
input3 = st.text_input("The final language for translating the text") |
|
|
|
if st.button("Generate"): |
|
result = process_input(input1, input2, input3) |
|
st.markdown(f"<div>{result}</div>", unsafe_allow_html = True) |
|
|
|
def process_input(a, b, c): |
|
model = ChatOpenAI(model = "gpt-4o-mini") |
|
|
|
|
|
prompt_one = PromptTemplate( |
|
input_variables = ["object"], |
|
template = "A story about {object}", |
|
) |
|
|
|
chain_one = LLMChain( |
|
llm = model, |
|
prompt = prompt_one, |
|
output_key = "story" |
|
) |
|
|
|
|
|
prompt_two = PromptTemplate( |
|
input_variables = ["style", "story"], |
|
template = "You are a professional writer. You will be given a short story and you will inspire from this original story and create a new story in {style} style. The original story: {story}.", |
|
) |
|
|
|
chain_two = LLMChain( |
|
llm = model, |
|
prompt = prompt_two, |
|
output_key = "expanded_story" |
|
|
|
) |
|
|
|
|
|
prompt_three = PromptTemplate( |
|
input_variables = ["expanded_story", "language"], |
|
template = "Translate the story: {expanded_story} into {language}?", |
|
) |
|
|
|
chain_three = LLMChain( |
|
llm = model, |
|
prompt = prompt_three, |
|
output_key = "translated" |
|
|
|
) |
|
|
|
|
|
final_chain = SequentialChain( |
|
chains = [chain_one, chain_two, chain_three], |
|
input_variables = ["object", "style", "language"], |
|
output_variables = ["translated"], |
|
verbose = True |
|
) |
|
|
|
resp = final_chain({ |
|
"object": a, |
|
"style":b, |
|
"language": c |
|
}) |
|
|
|
return resp["translated"] |
|
|
|
if __name__ == "__main__": |
|
main() |