Spaces:
Sleeping
Sleeping
| # Author: Brian King | |
| # For: BrandMuscle, Copyright 2023 All Rights Reserved | |
| import streamlit as st | |
| import os | |
| from llama_index import ( | |
| ServiceContext, | |
| SimpleDirectoryReader, | |
| VectorStoreIndex, | |
| ) | |
| from llama_index.llms import OpenAI | |
| import openai | |
| # Define Streamlit layout and interaction | |
| st.title("Streamlit App for PDF Retrieval and Text Generation") | |
| # Upload PDF | |
| uploaded_file = st.file_uploader("Choose a PDF file", type="pdf") | |
| def load_data(uploaded_file): | |
| with st.spinner('Indexing document...'): | |
| # Save the uploaded file temporarily | |
| with open("temp.pdf", "wb") as f: | |
| f.write(uploaded_file.read()) | |
| # Read and index documents using SimpleDirectoryReader | |
| reader = SimpleDirectoryReader(input_dir="./", recursive=False) | |
| docs = reader.load_data() | |
| service_context = ServiceContext.from_defaults( | |
| llm=OpenAI( | |
| model="gpt-3.5-turbo-16k", | |
| temperature=0.1, | |
| ), | |
| system_prompt="You are an AI assistant that uses context from a PDF to assist the user in generating text." | |
| ) | |
| index = VectorStoreIndex.from_documents(docs, service_context=service_context) | |
| return index | |
| # Placeholder for document indexing | |
| if uploaded_file is not None: | |
| index = load_data(uploaded_file) | |
| # Take user query input | |
| user_query = st.text_input("Search for the products/info you want to use to ground your generated text content:") | |
| # Initialize session_state for retrieved_text if not already present | |
| if 'retrieved_text' not in st.session_state: | |
| st.session_state['retrieved_text'] = '' | |
| # Search and display retrieved text | |
| if st.button("Retrieve"): | |
| with st.spinner('Retrieving text...'): | |
| # Use VectorStoreIndex to search | |
| query_engine = index.as_query_engine(similarity_top_k=3) | |
| st.session_state['retrieved_text'] = query_engine.query(user_query) | |
| st.write(f"Retrieved Text: {st.session_state['retrieved_text']}") | |
| # Select content type | |
| content_type = st.selectbox("Select content type:", ["Blog", "Tweet"]) | |
| # Generate text based on retrieved text and selected content type | |
| if st.button("Generate") and content_type: | |
| with st.spinner('Generating text...'): | |
| # Generate text using OpenAI API | |
| openai.api_key = os.getenv("OPENAI_API_KEY") | |
| try: | |
| if content_type == "Blog": | |
| prompt = f"Write a blog about 500 words in length using the {st.session_state['retrieved_text']}" | |
| elif content_type == "Tweet": | |
| prompt = f"Compose a tweet using the {st.session_state['retrieved_text']}" | |
| response = openai.ChatCompletion.create( | |
| model="gpt-3.5-turbo-16k", | |
| messages=[ | |
| {"role": "system", "content": "You are a helpful assistant."}, | |
| {"role": "user", "content": prompt} | |
| ] | |
| ) | |
| generated_text = response['choices'][0]['message']['content'] | |
| st.write(f"Generated Text: {generated_text}") | |
| except Exception as e: | |
| st.write(f"An error occurred: {e}") | |