Update app.py

app.py

@@ -2,28 +2,21 @@
 # import os
 # import re
 # import torch
-# from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
+# from transformers import AutoModelForCausalLM, AutoTokenizer
 # from PyPDF2 import PdfReader
 # from peft import get_peft_model, LoraConfig, TaskType
 
-# # ✅ Fix CUDA Memory Fragmentation
-# os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True"
+# # ✅ Force CPU execution for Streamlit Cloud
+# device = torch.device("cpu")
 
-# # 🔹 Load IBM Granite Model with 4-bit Quantization
+# # 🔹 Load IBM Granite Model (CPU-Compatible)
 # MODEL_NAME = "ibm-granite/granite-3.1-2b-instruct"
-# quant_config = BitsAndBytesConfig(load_in_4bit=True)  # Use 4-bit quantization
-
-# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-# # ✅ Ensure model initialization correctly
-# torch.cuda.empty_cache()  # Clear GPU memory before loading model
 
 # model = AutoModelForCausalLM.from_pretrained(
 #     MODEL_NAME,
-#     quantization_config=quant_config,
-#     device_map="auto",  # Auto-assign layers to available GPUs/CPUs
-#     torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32  # Use FP16 if GPU is available
-# ).to(device)  # Move model to correct device
+#     device_map="cpu",  # Force CPU execution
+#     torch_dtype=torch.float32  # Use float32 since Streamlit runs on CPU
+# )
 
 # tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
@@ -51,7 +44,6 @@
     
 #     return file_context.strip()
 
-# # 🛠 Function to Format AI Prompts
 # # 🛠 Function to Format AI Prompts
 # def format_prompt(system_msg, user_msg, file_context=""):
 #     if file_context:
@@ -60,10 +52,9 @@
 #         {"role": "system", "content": system_msg},
 #         {"role": "user", "content": user_msg}
 #     ]
+
 # # 🛠 Function to Generate AI Responses
 # def generate_response(input_text, max_tokens=1000, top_p=0.9, temperature=0.7):
-#     torch.cuda.empty_cache()  # ✅ Clear GPU memory before inference
-    
 #     model_inputs = tokenizer([input_text], return_tensors="pt").to(device)
     
 #     with torch.no_grad():
@@ -125,17 +116,6 @@
 #         # 🔹 User Input for Analysis
 #         user_prompt = "Perform a detailed technical analysis of the attached contract document, highlighting potential risks, legal pitfalls, compliance issues, and areas where contractual terms may lead to future disputes or operational challenges."
 
-#         # user_prompt = st.text_area(
-#         #     "📝 Describe what you want to analyze:",
-#         #     "Perform a detailed technical analysis of the attached contract document, highlighting potential risks, legal pitfalls, compliance issues, and areas where contractual terms may lead to future disputes or operational challenges."
-#         # )
-#         # with st.empty():  # This hides the text area
-#         #   user_prompt = st.text_area(
-#         #       "📝 Describe what you want to analyze:",
-#         #       "Perform a detailed technical analysis of the attached contract document, highlighting potential risks, legal pitfalls, compliance issues, and areas where contractual terms may lead to future disputes or operational challenges."
-#         #   )
-
-
 #         if st.button("🔍 Analyze Document"):
 #             with st.spinner("Analyzing contract document... ⏳"):
 #                 final_answer = granite_simple(user_prompt, temp_file_path)
@@ -154,6 +134,7 @@
 
 
 
+
 import streamlit as st
 import os
 import re
@@ -162,16 +143,16 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 from PyPDF2 import PdfReader
 from peft import get_peft_model, LoraConfig, TaskType
 
-# ✅ Force CPU execution for Streamlit Cloud
-device = torch.device("cpu")
+# ✅ Auto-detect GPU for Hugging Face Spaces
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-# 🔹 Load IBM Granite Model (CPU-Compatible)
+# 🔹 Load IBM Granite Model (CPU/GPU Compatible)
 MODEL_NAME = "ibm-granite/granite-3.1-2b-instruct"
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
-    device_map="cpu",  # Force CPU execution
-    torch_dtype=torch.float32  # Use float32 since Streamlit runs on CPU
+    device_map="auto",  # Auto-detect GPU if available
+    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32
 )
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
@@ -188,7 +169,7 @@ lora_config = LoraConfig(
 model = get_peft_model(model, lora_config)
 model.eval()
 
-# 🛠 Function to Read & Extract Text from PDFs
+# 🛠 Function to Read & Extract Text from PDFs (No Temp File Needed)
 def read_files(file):
     file_context = ""
     reader = PdfReader(file)
@@ -233,13 +214,11 @@ def post_process(text):
     unique_lines = list(dict.fromkeys([line.strip() for line in lines if line.strip()]))
     return "\n".join(unique_lines)
 
-# 🛠 Function to Handle RAG with IBM Granite & Streamlit
-def granite_simple(prompt, file):
-    file_context = read_files(file) if file else ""
-    
+# 🛠 Function to Handle AI Analysis (No Temp File)
+def granite_simple(prompt, file_content):
     system_message = "You are IBM Granite, a legal AI assistant specializing in contract analysis."
     
-    messages = format_prompt(system_message, prompt, file_context)
+    messages = format_prompt(system_message, prompt, file_content)
     input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
     
     response = generate_response(input_text)
@@ -259,30 +238,26 @@ def main():
         top_p = st.slider("Top P (sampling)", 0.1, 1.0, 0.9, 0.1)
         temperature = st.slider("Temperature (creativity)", 0.1, 1.0, 0.7, 0.1)
 
-    # 🔹 File Upload Section
+    # 🔹 File Upload Section (No Temp File)
     uploaded_file = st.file_uploader("📂 Upload a contract document (PDF)", type="pdf")
 
     if uploaded_file is not None:
-        temp_file_path = "temp_uploaded_contract.pdf"
-        with open(temp_file_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-
         st.success("✅ File uploaded successfully!")
 
+        # 🔹 Read PDF Content (No Temp File)
+        file_content = read_files(uploaded_file)
+
         # 🔹 User Input for Analysis
         user_prompt = "Perform a detailed technical analysis of the attached contract document, highlighting potential risks, legal pitfalls, compliance issues, and areas where contractual terms may lead to future disputes or operational challenges."
 
         if st.button("🔍 Analyze Document"):
             with st.spinner("Analyzing contract document... ⏳"):
-                final_answer = granite_simple(user_prompt, temp_file_path)
+                final_answer = granite_simple(user_prompt, file_content)
 
             # 🔹 Display Analysis Result
             st.subheader("📑 Analysis Result")
             st.write(final_answer)
 
-            # 🔹 Remove Temporary File
-            os.remove(temp_file_path)
-
 # 🔥 Run Streamlit App
 if __name__ == '__main__':
-    main()
+    main()