File size: 2,195 Bytes
226100d
 
5cefde2
226100d
 
 
 
5cefde2
226100d
 
8a64540
e5a7929
a2bb924
 
 
 
724b7ed
 
 
a2bb924
724b7ed
 
 
 
 
 
 
 
a2bb924
724b7ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2d065df
724b7ed
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
e5a7929
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import os
import streamlit as st  # Front End / UI
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")
    
    # Chain 1
    prompt_one = PromptTemplate(
        input_variables = ["object"],
        template = "A story about {object}",
    )

    chain_one = LLMChain(
        llm = model, 
        prompt = prompt_one,
        output_key = "story"
    )

    # Chain 2
    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"
        
    )

    # Chain 3
    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"
        
    )

    # SEQUENCE
    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()