Update app.py

app.py

@@ -20,7 +20,7 @@ HF_TOKEN = os.getenv('HF_TOKEN')
 AZURE_TRANSLATION_KEY = os.getenv('AZURE_TRANSLATION_KEY')
 
 class Translator:
-    def _init_(self):
+    def __init__(self):  # Fixed method name from _init_ to __init__
         self.key = AZURE_TRANSLATION_KEY
         self.region = 'centralindia'
         self.endpoint = "https://api.cognitive.microsofttranslator.com"
@@ -99,7 +99,7 @@ class TextExtractor:
         return "Unsupported input type"
 
 class LegalEaseAssistant:
-    def _init_(self):
+    def __init__(self):  # Fixed method name from _init_ to __init__
         if not HF_TOKEN:
             raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
         
@@ -135,7 +135,7 @@ class LegalEaseAssistant:
         
         prompt = task_prompts.get(task_type, f"Analyze the following text and provide points:\n\n{text}\n\nAnalysis:")
         
-        inputs = self.tokenizer(prompt, return_tensors="pt")
+        inputs = self.tokenizer(prompt, return_tensors="pt").to(self.model.device)  # Add to(device) to ensure tensor is on the right device
         outputs = self.model.generate(
             **inputs, 
             max_new_tokens=300,
@@ -207,7 +207,7 @@ def create_interface():
                 with gr.Row():
                     with gr.Column(scale=1):
                         simplify_input = gr.File(
-                            file_types=['txt', 'pdf', 'image'], 
+                            file_types=['txt', 'pdf', 'jpg', 'jpeg', 'png'],  # Added image file types explicitly
                             label="📎 Upload Document"
                         )
                         gr.HTML("<div style='height: 10px'></div>")
@@ -250,7 +250,7 @@ def create_interface():
                 with gr.Row():
                     with gr.Column(scale=1):
                         summary_input = gr.File(
-                            file_types=['txt', 'pdf', 'image'], 
+                            file_types=['txt', 'pdf', 'jpg', 'jpeg', 'png'],  # Added image file types explicitly
                             label="📎 Upload Document"
                         )
                         gr.HTML("<div style='height: 10px'></div>")
@@ -294,7 +294,7 @@ def create_interface():
                 with gr.Row():
                     with gr.Column(scale=1):
                         terms_input = gr.File(
-                            file_types=['txt', 'pdf', 'image'], 
+                            file_types=['txt', 'pdf', 'jpg', 'jpeg', 'png'],  # Added image file types explicitly
                             label="📎 Upload Document"
                         )
                         gr.HTML("<div style='height: 10px'></div>")
@@ -338,7 +338,7 @@ def create_interface():
                 with gr.Row():
                     with gr.Column(scale=1):
                         contract1_input = gr.File(
-                            file_types=['txt', 'pdf', 'image'], 
+                            file_types=['txt', 'pdf', 'jpg', 'jpeg', 'png'],  # Added image file types explicitly
                             label="📎 Upload First Contract"
                         )
                         gr.HTML("<div style='height: 10px'></div>")
@@ -350,7 +350,7 @@ def create_interface():
                     
                     with gr.Column(scale=1):
                         contract2_input = gr.File(
-                            file_types=['txt', 'pdf', 'image'], 
+                            file_types=['txt', 'pdf', 'jpg', 'jpeg', 'png'],  # Added image file types explicitly
                             label="📎 Upload Second Contract"
                         )
                         gr.HTML("<div style='height: 10px'></div>")
@@ -378,22 +378,25 @@ def create_interface():
                     if not contract1 or not contract2:
                         return "Please provide both contracts for comparison."
                     
+                    # Extract text if needed
+                    if not isinstance(contract1, str):
+                        contract1 = assistant.text_extractor.extract_text_from_input(contract1)
+                    if not isinstance(contract2, str):
+                        contract2 = assistant.text_extractor.extract_text_from_input(contract2)
+                    
                     def compare_contracts(contract1, contract2):
                         prompt = f"""Compare these two contracts and identify key differences and similarities:
-
 Contract 1:
 {contract1}
-
 Contract 2:
 {contract2}
-
 Please analyze and list:
 1. Key similarities
 2. Important differences
 3. Unique terms in each contract
 4. Potential implications of the differences"""
                         
-                        inputs = assistant.tokenizer(prompt, return_tensors="pt")
+                        inputs = assistant.tokenizer(prompt, return_tensors="pt").to(assistant.model.device)  # Add to(device) to ensure tensor is on the right device
                         outputs = assistant.model.generate(
                             **inputs,
                             max_new_tokens=400,
@@ -427,7 +430,7 @@ Please analyze and list:
                 with gr.Row():
                     with gr.Column(scale=1):
                         risk_input = gr.File(
-                            file_types=['txt', 'pdf', 'image'], 
+                            file_types=['txt', 'pdf', 'jpg', 'jpeg', 'png'],  # Added image file types explicitly
                             label="📎 Upload Document"
                         )
                         gr.HTML("<div style='height: 10px'></div>")