Spaces:
Running
Running
File size: 3,585 Bytes
b9fb8a5 aa1b98f b9fb8a5 aa1b98f 491ad49 aa1b98f 04a08d3 aa1b98f 04a08d3 aa1b98f |
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 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
import streamlit as st
import openai
import anthropic
import fitz
# Set up API clients
openai_api_key = st.secrets["OPENAI_API_KEY"]
openai_client = openai.OpenAI(api_key=openai_api_key)
anthropic_api_key = st.secrets["ANTHROPIC_API_KEY"]
anthropic_client = anthropic.Anthropic(api_key=anthropic_api_key)
# Password protection
def check_password():
if "password_correct" not in st.session_state:
st.session_state["password_correct"] = False
if not st.session_state["password_correct"]:
password = st.text_input("Enter the password", type="password")
if st.button("Submit"):
if password == st.secrets["PASSWORD"]:
st.session_state["password_correct"] = True
st.rerun()
else:
st.error("😕 Password incorrect")
return False
else:
return True
def extract_text_from_pdf(file):
try:
with fitz.open(stream=file.read(), filetype="pdf") as doc:
text = ""
for page in doc:
text += page.get_text()
return text
except Exception as e:
st.error(f"Error parsing PDF: {str(e)}")
return None
def main():
st.title("Harvey Legal Research Take-Home")
if not check_password():
return
# Model selector
model = st.selectbox("Select Model", ["GPT-4o", "Claude 3.5 Sonnet"])
# Prompt input
prompt = st.text_area("Enter your prompt:", height=200)
# Document upload
uploaded_file = st.file_uploader("Upload a document", type=["txt", "pdf"])
document_text = None
if uploaded_file is not None:
if uploaded_file.type == "text/plain":
document_text = uploaded_file.getvalue().decode()
elif uploaded_file.type == "application/pdf":
document_text = extract_text_from_pdf(uploaded_file)
else:
st.error("Unsupported file type. Please upload a text or PDF file.")
# Run button
if st.button("Run"):
if not prompt:
st.warning("Please enter a prompt.")
return
if document_text is None:
full_prompt = prompt
else:
full_prompt = f"Document content:\n\n{document_text}\n\nPrompt:\n\n{prompt}"
if model == "GPT-4o":
with st.spinner("Processing with GPT-4o..."):
stream = openai_client.chat.completions.create(
model="gpt-4o",
messages=[{"role": "user", "content": full_prompt}],
stream=True,
)
full_response = ""
response_area = st.empty()
for chunk in stream:
if chunk.choices[0].delta.content is not None:
full_response += chunk.choices[0].delta.content
response_area.markdown(f"**Response:**\n\n{full_response}")
elif model == "Claude 3.5 Sonnet":
with st.spinner("Processing with Claude 3.5 Sonnet..."):
with anthropic_client.messages.stream(
model="claude-3-5-sonnet-20240620",
max_tokens=8192,
messages=[{"role": "user", "content": full_prompt}],
) as stream:
full_response = ""
response_area = st.empty()
for text in stream.text_stream:
full_response += text
response_area.markdown(f"**Response:**\n\n{full_response}")
if __name__ == "__main__":
main()
|