Spaces:
Sleeping
Sleeping
| import os | |
| import streamlit as st | |
| from PyPDF2 import PdfReader | |
| from langchain.text_splitter import RecursiveCharacterTextSplitter | |
| from langchain.embeddings.openai import OpenAIEmbeddings | |
| from langchain.vectorstores import FAISS | |
| from langchain.chains.question_answering import load_qa_chain | |
| from langchain.llms import OpenAI | |
| from reportlab.lib.pagesizes import letter | |
| from reportlab.pdfgen import canvas | |
| from io import BytesIO | |
| from deep_translator import GoogleTranslator | |
| # Set your OpenAI API key here | |
| os.environ["OPENAI_API_KEY"] = os.getenv("OpenAIKey") | |
| st.title("LegalAIDe") | |
| if 'translated_answer' not in st.session_state: | |
| st.session_state.translated_answer = "" | |
| uploaded_file = st.file_uploader("Choose a PDF file", type="pdf") | |
| if uploaded_file is not None: | |
| pdf_reader = PdfReader(uploaded_file) | |
| raw_text = '' | |
| for page in pdf_reader.pages: | |
| text = page.extract_text() | |
| if text: | |
| raw_text += text | |
| text_splitter = RecursiveCharacterTextSplitter( | |
| separators=["\n\n", "\n", ".", "!", "?"], | |
| chunk_size=1000, | |
| chunk_overlap=200, | |
| length_function=len | |
| ) | |
| texts = text_splitter.split_text(raw_text) | |
| embeddings = OpenAIEmbeddings() | |
| document_search = FAISS.from_texts(texts, embeddings) | |
| query = st.text_input("Enter your question:") | |
| if st.button("Get Answer"): | |
| docs = document_search.similarity_search(query) | |
| chain = load_qa_chain(OpenAI(), chain_type="stuff") | |
| # Add context about your role | |
| context = "You are a lawyer and need assistance with legal questions." | |
| # Use a more explicit prompt for better context | |
| prompt = f"Context: {context}\n\nQuestion: {query}\n\nPlease provide a detailed answer based on the given documents." | |
| answer = chain.run(input_documents=docs, question=prompt) | |
| st.write("Answer:", answer) | |
| st.session_state.answer = answer | |
| # Language selection and translation | |
| translation_language = st.text_input("Enter translation language:") # Corrected to use st.text_input | |
| if st.button("Translate Answer"): | |
| if translation_language: # Check if a language was entered | |
| try: | |
| translated_answer = GoogleTranslator(source='auto', target=translation_language).translate(st.session_state.answer) | |
| st.session_state.translated_answer = translated_answer | |
| except Exception as e: | |
| st.error(f"An error occurred: {e}") | |
| else: | |
| st.error("Please enter a valid language.") | |
| # Display translated answer if available | |
| if st.session_state.get('translated_answer'): | |
| st.write("Translated Answer:", st.session_state.translated_answer) | |
| if st.button("Generate PDF"): | |
| # Generate PDF | |
| pdf_buffer = BytesIO() | |
| c = canvas.Canvas(pdf_buffer, pagesize=letter) | |
| c.drawString(100, 750, "Question:") | |
| c.drawString(100, 735, query) | |
| c.drawString(100, 715, "Answer:") | |
| text_lines = st.session_state.answer.split('\n') | |
| y = 700 | |
| for line in text_lines: | |
| if y < 50: | |
| c.showPage() | |
| y = 750 | |
| c.drawString(100, y, line) | |
| y -= 15 | |
| if st.session_state.translated_answer: | |
| c.drawString(100, y, "Translated Answer:") | |
| y -= 15 | |
| text_lines = st.session_state.translated_answer.split('\n') | |
| for line in text_lines: | |
| if y < 50: | |
| c.showPage() | |
| y = 750 | |
| c.drawString(100, y, line) | |
| y -= 15 | |
| c.save() | |
| pdf_buffer.seek(0) | |
| st.download_button( | |
| label="Download Answer as PDF", | |
| data=pdf_buffer, | |
| file_name="answer.pdf", | |
| mime="application/pdf" | |
| ) | |