brianjking commited on
Commit
8bcc00e
1 Parent(s): 71f9869

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +81 -0
app.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Author: Brian King
2
+ # For: BrandMuscle, Copyright 2023 All Rights Reserved
3
+
4
+ import streamlit as st
5
+ import os
6
+ from llama_index import (
7
+ ServiceContext,
8
+ SimpleDirectoryReader,
9
+ VectorStoreIndex,
10
+ )
11
+ from llama_index.llms import OpenAI
12
+ import openai
13
+
14
+ # Define Streamlit layout and interaction
15
+ st.title("Streamlit App for PDF Retrieval and Text Generation")
16
+
17
+ # Upload PDF
18
+ uploaded_file = st.file_uploader("Choose a PDF file", type="pdf")
19
+
20
+ @st.cache_resource(show_spinner=False)
21
+ def load_data(uploaded_file):
22
+ with st.spinner('Indexing document...'):
23
+ # Save the uploaded file temporarily
24
+ with open("temp.pdf", "wb") as f:
25
+ f.write(uploaded_file.read())
26
+ # Read and index documents using SimpleDirectoryReader
27
+ reader = SimpleDirectoryReader(input_dir="./", recursive=False)
28
+ docs = reader.load_data()
29
+ service_context = ServiceContext.from_defaults(
30
+ llm=OpenAI(
31
+ model="gpt-3.5-turbo-16k",
32
+ temperature=0.1,
33
+ ),
34
+ system_prompt="You are an AI assistant that uses context from a PDF to assist the user in generating text."
35
+ )
36
+ index = VectorStoreIndex.from_documents(docs, service_context=service_context)
37
+ return index
38
+
39
+ # Placeholder for document indexing
40
+ if uploaded_file is not None:
41
+ index = load_data(uploaded_file)
42
+
43
+ # Take user query input
44
+ user_query = st.text_input("Search for the products/info you want to use to ground your generated text content:")
45
+
46
+ # Initialize session_state for retrieved_text if not already present
47
+ if 'retrieved_text' not in st.session_state:
48
+ st.session_state['retrieved_text'] = ''
49
+
50
+ # Search and display retrieved text
51
+ if st.button("Retrieve"):
52
+ with st.spinner('Retrieving text...'):
53
+ # Use VectorStoreIndex to search
54
+ query_engine = index.as_query_engine(similarity_top_k=3)
55
+ st.session_state['retrieved_text'] = query_engine.query(user_query)
56
+ st.write(f"Retrieved Text: {st.session_state['retrieved_text']}")
57
+
58
+ # Select content type
59
+ content_type = st.selectbox("Select content type:", ["Blog", "Tweet"])
60
+
61
+ # Generate text based on retrieved text and selected content type
62
+ if st.button("Generate") and content_type:
63
+ with st.spinner('Generating text...'):
64
+ # Generate text using OpenAI API
65
+ openai.api_key = os.getenv("OPENAI_API_KEY")
66
+ try:
67
+ if content_type == "Blog":
68
+ prompt = f"Write a blog about 500 words in length using the {st.session_state['retrieved_text']}"
69
+ elif content_type == "Tweet":
70
+ prompt = f"Compose a tweet using the {st.session_state['retrieved_text']}"
71
+ response = openai.ChatCompletion.create(
72
+ model="gpt-3.5-turbo-16k",
73
+ messages=[
74
+ {"role": "system", "content": "You are a helpful assistant."},
75
+ {"role": "user", "content": prompt}
76
+ ]
77
+ )
78
+ generated_text = response['choices'][0]['message']['content']
79
+ st.write(f"Generated Text: {generated_text}")
80
+ except Exception as e:
81
+ st.write(f"An error occurred: {e}")