astroarya's picture
Update app.py
90a5b2d verified
raw
history blame
1.88 kB
import os
import streamlit as st
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain.utilities import WikipediaAPIWrapper
from transformers import GPT2LMHeadModel, GPT2Tokenizer
# Load the GPT-2 model and tokenizer
model_name = "gpt2"
tokenizer = GPT2Tokenizer.from_pretrained(model_name)
model = GPT2LMHeadModel.from_pretrained(model_name)
# app framework
st.title('Youtube Content Studio πŸš€')
prompt = st.text_input('Plug in your prompt here')
# Prompt templates
title_template = PromptTemplate(
input_variables=['topic'],
template='write me a youtube video title about {topic}'
)
script_template = PromptTemplate(
input_variables=['title', 'wikipedia_research'],
template='write me a youtube script based on this title TITLE: {title} while leveraging this wikipedia research: {wikipedia_research}'
)
# Memory
title_memory = ConversationBufferMemory(input_key='topic', memory_key='chat_history')
script_memory = ConversationBufferMemory(input_key='title', memory_key='chat_history')
# LLMS
llm = model
title_chain = LLMChain(llm=llm, prompt=title_template, verbose=True, output_key='title', memory=title_memory)
script_chain = LLMChain(llm=llm, prompt=script_template, verbose=True, output_key='script', memory=script_memory)
wiki = WikipediaAPIWrapper()
# Show stuff to the screen if there is a prompt
if prompt:
title = title_chain.run(prompt)
wiki_research = wiki.run(prompt)
script = script_chain.run(title=title, wikipedia_research=wiki_research)
st.write(title)
st.write(script)
with st.expander('Title History'):
st.info(title_memory.buffer)
with st.expander('Script History'):
st.info(script_memory.buffer)
with st.expander('Wikipedia research History'):
st.info(wiki_research)