Update app.py

app.py

@@ -1,32 +1,33 @@
 import streamlit as st
 import pandas as pd
-import os
 from transformers import T5ForConditionalGeneration, T5Tokenizer
-import groq
+from docx import Document
 
-# Initialize Groq API
-groq_client = groq.Client(api_key="financialstatements")
-
-# Load RAG components
-retriever_tokenizer = DPRContextEncoderTokenizer.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
-retriever_model = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base")
+# Load the generator model (FLAN-T5)
 generator_tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
 generator_model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
 
+# Function to read DOCS files
+def read_docs(file):
+    doc = Document(file)
+    text = "\n".join([paragraph.text for paragraph in doc.paragraphs])
+    return text
+
 # Function to process user input and generate financial statements
 def generate_financial_statements(file, file_type):
-    # Read the file
+    # Read the file based on its type
     if file_type == "csv":
         df = pd.read_csv(file)
+        context = df.to_string()
     elif file_type == "excel":
         df = pd.read_excel(file)
+        context = df.to_string()
+    elif file_type == "docs":
+        context = read_docs(file)
     else:
-        st.error("Unsupported file type. Please upload a CSV or Excel file.")
+        st.error("Unsupported file type. Please upload a CSV, Excel, or DOCS file.")
         return
 
-    # Convert the data into a context string
-    context = df.to_string()
-
     # Define financial statement queries
     queries = [
         "Generate a journal from the following financial data:",
@@ -36,19 +37,15 @@ def generate_financial_statements(file, file_type):
         "Generate a cash flow statement from the following financial data:"
     ]
 
-    # Generate financial statements using RAG
+    # Generate financial statements using the generator model
     financial_statements = {}
     for query in queries:
         # Combine query and context
         input_text = f"{query}\n{context}"
 
-        # Retrieve relevant information (optional, if using a retriever)
-        input_ids = retriever_tokenizer(input_text, return_tensors="pt").input_ids
-        retrieved_context = retriever_model(input_ids)
-
         # Generate response using the generator model
-        input_ids = generator_tokenizer(input_text, return_tensors="pt").input_ids
-        output = generator_model.generate(input_ids)
+        input_ids = generator_tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True).input_ids
+        output = generator_model.generate(input_ids, max_length=512)
         response = generator_tokenizer.decode(output[0], skip_special_tokens=True)
 
         # Store the result
@@ -58,10 +55,10 @@ def generate_financial_statements(file, file_type):
 
 # Streamlit UI
 st.title("Financial Statement Generator")
-st.write("Upload your financial data (CSV or Excel) to generate journal, general ledger, income statement, balance sheet, and cash flow statement.")
+st.write("Upload your financial data (CSV, Excel, or DOCS) to generate journal, general ledger, income statement, balance sheet, and cash flow statement.")
 
 # File upload
-uploaded_file = st.file_uploader("Upload your file", type=["csv", "xlsx"])
+uploaded_file = st.file_uploader("Upload your file", type=["csv", "xlsx", "docx"])
 if uploaded_file is not None:
     file_type = uploaded_file.name.split(".")[-1]
     financial_statements = generate_financial_statements(uploaded_file, file_type)