Update app.py

app.py

@@ -6,21 +6,48 @@ from transformers import AutoTokenizer, AutoModelForCausalLM
 from huggingface_hub import login
 import pytesseract
 from PIL import Image
-import PyPDF2
+import fitz  # PyMuPDF
 import requests
 import uuid
 
 # Configuration
 MODEL_NAME = "google/gemma-2b-it"
 CURRENT_USER = "AkarshanGupta"
-CURRENT_TIME = "2025-03-22 21:00:45"
+CURRENT_TIME = "2025-03-23 03:33:01"
 
 # API Keys
 HF_TOKEN = os.getenv('HF_TOKEN')
 AZURE_TRANSLATION_KEY = os.getenv('AZURE_TRANSLATION_KEY')
+LLAMA_API_KEY = os.getenv('LLAMA_API_KEY')
+LLAMA_API_ENDPOINT = "https://api.llama.ai/v1/generate"
+
+class TextExtractor:
+    @staticmethod
+    def extract_text_from_input(input_file):
+        if isinstance(input_file, str):
+            return input_file
+        
+        if isinstance(input_file, Image.Image):
+            try:
+                return pytesseract.image_to_string(input_file)
+            except Exception as e:
+                return f"Error extracting text from image: {str(e)}"
+        
+        if hasattr(input_file, 'name') and input_file.name.lower().endswith('.pdf'):
+            try:
+                doc = fitz.open(stream=input_file.read(), filetype="pdf")
+                text = ""
+                for page in doc:
+                    text += page.get_text() + "\n\n"
+                doc.close()
+                return text
+            except Exception as e:
+                return f"Error extracting text from PDF: {str(e)}"
+        
+        return "Unsupported input type"
 
 class Translator:
-    def init(self):
+    def _init_(self):
         self.key = AZURE_TRANSLATION_KEY
         self.region = 'centralindia'
         self.endpoint = "https://api.cognitive.microsofttranslator.com"
@@ -30,13 +57,11 @@ class Translator:
 
     def translate_text(self, text, target_language="en"):
         try:
-            # Split the text into bullet points
             bullet_points = text.split('\n• ')
             translated_points = []
             
-            # Translate each bullet point separately
             for point in bullet_points:
-                if point.strip():  # Only translate non-empty points
+                if point.strip():
                     path = '/translate'
                     constructed_url = self.endpoint + path
                     
@@ -47,7 +72,7 @@ class Translator:
                     
                     headers = {
                         'Ocp-Apim-Subscription-Key': self.key,
-                        'Ocp-Apim-Subscription-Region': 'centralindia',
+                        'Ocp-Apim-Subscription-Region': self.region,
                         'Content-type': 'application/json',
                         'X-ClientTraceId': str(uuid.uuid4())
                     }
@@ -67,44 +92,22 @@ class Translator:
                     translation = response.json()[0]["translations"][0]["text"]
                     translated_points.append(translation)
             
-            # Reconstruct the bullet-pointed text
             translated_text = '\n• ' + '\n• '.join(translated_points)
             return translated_text
             
         except Exception as e:
             return f"Translation error: {str(e)}"
 
-class TextExtractor:
-    @staticmethod
-    def extract_text_from_input(input_file):
-        if isinstance(input_file, str):
-            return input_file
-        
-        if isinstance(input_file, Image.Image):
-            try:
-                return pytesseract.image_to_string(input_file)
-            except Exception as e:
-                return f"Error extracting text from image: {str(e)}"
-        
-        if hasattr(input_file, 'name') and input_file.name.lower().endswith('.pdf'):
-            try:
-                pdf_reader = PyPDF2.PdfReader(input_file)
-                text = ""
-                for page in pdf_reader.pages:
-                    text += page.extract_text() + "\n\n"
-                return text
-            except Exception as e:
-                return f"Error extracting text from PDF: {str(e)}"
-        
-        return "Unsupported input type"
-
 class LegalEaseAssistant:
-    def init(self):
+    def _init_(self):
         if not HF_TOKEN:
             raise ValueError("Hugging Face token not found. Please set the HF_TOKEN environment variable.")
         
         login(token=HF_TOKEN)
         
+        # Initialize text_extractor first
+        self.text_extractor = TextExtractor()
+        
         self.tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME, 
             token=HF_TOKEN
@@ -115,10 +118,8 @@ class LegalEaseAssistant:
             device_map="auto",
             torch_dtype=torch.float32
         )
-        self.text_extractor = TextExtractor()
     
     def format_response(self, text):
-        """Format response as bullet points"""
         sentences = [s.strip() for s in text.split('.') if s.strip()]
         bullet_points = ['• ' + s + '.' for s in sentences]
         return '\n'.join(bullet_points)
@@ -150,6 +151,23 @@ class LegalEaseAssistant:
         raw_response = response_parts[-1].strip() if len(response_parts) > 1 else response.strip()
         
         return self.format_response(raw_response)
+    
+    def generate_chatbot_response(self, user_input):
+        if not LLAMA_API_KEY:
+            return "LLaMA API key not found. Please set the LLAMA_API_KEY environment variable."
+        
+        response = requests.post(
+            LLAMA_API_ENDPOINT,
+            headers={"Authorization": f"Bearer {LLAMA_API_KEY}"},
+            json={"prompt": user_input, "max_tokens": 150}
+        )
+        
+        if response.status_code == 401:
+            return "Unauthorized: Please check your LLaMA API key."
+        elif response.status_code != 200:
+            return f"Error: Received {response.status_code} status code from LLaMA API."
+        
+        return response.json()["choices"][0]["text"].strip()
 
 def create_interface():
     assistant = LegalEaseAssistant()
@@ -179,16 +197,20 @@ def create_interface():
             result = translator.translate_text(result, SUPPORTED_LANGUAGES[target_lang])
         return result
 
-    with gr.Blocks(title="LegalEase: AI Legal Assistant") as demo:
+    with gr.Blocks(title="LegalEase", css="""
+        .gradio-container {max-width: 1200px; margin: auto;}
+        .header {text-align: center; margin-bottom: 2rem;}
+        .content {padding: 2rem;}
+    """) as demo:
         gr.HTML(f"""
-        <div style="text-align: center; background-color: #f0f2f6; padding: 20px; border-radius: 10px; margin-bottom: 20px;">
+        <div style="text-align: center; background-color: #e0e0e0; padding: 20px; border-radius: 10px; margin-bottom: 20px;">
             <h1 style="color: #2c3e50; font-size: 2.5em; margin-bottom: 10px;">📜 LegalEase</h1>
             <h2 style="color: #34495e; font-size: 1.5em; margin-bottom: 20px;">AI-Powered Legal Document Assistant</h2>
             <div style="display: flex; justify-content: center; gap: 40px; color: #576574; font-size: 1.1em;">
-                <div style="background-color: white; padding: 10px 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
+                <div style="background-color: #e0e0e0; padding: 10px 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                     <span style="font-weight: bold;">User:</span> {CURRENT_USER}
                 </div>
-                <div style="background-color: white; padding: 10px 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
+                <div style="background-color: #e0e0e0; padding: 10px 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
                     <span style="font-weight: bold;">Last Updated:</span> {CURRENT_TIME} UTC
                 </div>
             </div>
@@ -203,6 +225,7 @@ def create_interface():
         )
         
         with gr.Tabs():
+            # Simplify Language Tab
             with gr.Tab("📝 Simplify Language"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -246,6 +269,7 @@ def create_interface():
                     outputs=simplify_output
                 )
 
+            # Document Summary Tab
             with gr.Tab("📚 Document Summary"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -289,6 +313,51 @@ def create_interface():
                     outputs=summary_output
                 )
 
+            # Key Terms Tab
+            with gr.Tab("🔑 Key Terms"):
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        terms_input = gr.File(
+                            file_types=['txt', 'pdf', 'image'], 
+                            label="📎 Upload Document"
+                        )
+                        gr.HTML("<div style='height: 10px'></div>")
+                        terms_text_input = gr.Textbox(
+                            label="✍ Or Type/Paste Text",
+                            placeholder="Enter your legal document here...",
+                            lines=4
+                        )
+                        gr.HTML("<div style='height: 10px'></div>")
+                        terms_btn = gr.Button(
+                            "🔍 Extract Key Terms",
+                            variant="primary"
+                        )
+                    
+                    with gr.Column(scale=1):
+                        terms_output = gr.Textbox(
+                            label="🔑 Key Terms & Definitions",
+                            lines=12,
+                            show_copy_button=True
+                        )
+                
+                def terms_handler(file, text, lang):
+                    input_source = file or text
+                    if not input_source:
+                        return "Please provide some text or upload a document to analyze key terms."
+                    return process_with_translation(
+                        assistant.generate_response,
+                        input_source,
+                        "key_terms",
+                        target_lang=lang
+                    )
+                
+                terms_btn.click(
+                    fn=terms_handler,
+                    inputs=[terms_input, terms_text_input, language_selector],
+                    outputs=terms_output
+                )
+
+            # Risk Analysis Tab
             with gr.Tab("⚠ Risk Analysis"):
                 with gr.Row():
                     with gr.Column(scale=1):
@@ -332,16 +401,49 @@ def create_interface():
                     outputs=risk_output
                 )
 
+            # Legal Assistant Chat Tab
+            with gr.Tab("🤖 Legal Assistant Chat"):
+                chatbot_input = gr.Textbox(
+                    label="💬 Your Message",
+                    placeholder="Ask me anything about legal matters...",
+                    lines=2
+                )
+                chatbot_output = gr.Textbox(
+                    label="🤖 Assistant Response",
+                    lines=10,
+                    show_copy_button=True
+                )
+                chatbot_btn = gr.Button(
+                    "💬 Send Message",
+                    variant="primary"
+                )
+
+                def chatbot_handler(user_input, lang):
+                    if not user_input:
+                        return "Please type a message to start the conversation."
+                    response = assistant.generate_chatbot_response(user_input)
+                    if lang != "English":
+                        response = translator.translate_text(response, SUPPORTED_LANGUAGES[lang])
+                    return response
+                
+                chatbot_btn.click(
+                    fn=chatbot_handler,
+                    inputs=[chatbot_input, language_selector],
+                    outputs=chatbot_output
+                )
+
         gr.HTML(f"""
-        <div style="text-align: center; margin-top: 20px; padding: 20px; background-color: #f0f2f6; border-radius: 10px;">
-            <p style="color: #576574; margin: 0;">Powered by Gemma 2B and Azure Translator</p>
-            <p style="color: #576574; margin: 5px 0 0 0; font-size: 0.9em;">Built for Language Translation Hackathon</p>
+        <div style="text-align: center; margin-top: 20px; padding: 20px; background-color: #e0e0e0; border-radius: 10px;">
+            <p style="color: #576574; margin: 0;">Made by Team Ice Age</p>
         </div>
         """)
 
     return demo
 
-demo = create_interface()
+def main():
+    demo = create_interface()
+    demo.queue()
+    demo.launch(share=True)
 
-if _name_ == "_main_":
-    demo.launch()
+if __name__ == "__main__":
+    main()