Spaces:
Build error
Build error
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) | |