Update app.py

app.py

@@ -1,8 +1,132 @@
-import gradio as gr
-from ui import create_interface
-
-# Create the interface
-demo = create_interface()
-
-if __name__ == "__main__":
-    demo.launch()
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+import pytesseract
+from PIL import Image
+import PyPDF2
+import io
+import requests
+import os
+
+# Azure Translator API Configuration
+AZURE_TRANSLATOR_KEY = "your_azure_key"
+AZURE_TRANSLATOR_ENDPOINT = "https://api.cognitive.microsofttranslator.com/translate"
+AZURE_TRANSLATOR_REGION = "your_region"
+
+# Specify the model
+MODEL_NAME = "google/gemma-2b-it"
+
+class LegalEaseAssistant:
+    def __init__(self, model_name=MODEL_NAME):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, 
+            device_map="cpu",
+            load_in_8bit=True,
+            torch_dtype=torch.float16
+        )
+    
+    def extract_text_from_input(self, 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"
+    
+    def generate_response(self, input_file, task_type):
+        text = self.extract_text_from_input(input_file)
+        
+        task_prompts = {
+            "simplify": f"Simplify the following legal text:\n\n{text}\n\nSimplified explanation:",
+            "summary": f"Provide a concise summary:\n\n{text}\n\nSummary:",
+            "key_terms": f"Identify key legal terms:\n\n{text}\n\nKey Terms:",
+            "risk": f"Perform a risk analysis:\n\n{text}\n\nRisk Assessment:"
+        }
+        
+        prompt = task_prompts.get(task_type, f"Analyze the following text:\n\n{text}\n\nAnalysis:")
+        
+        inputs = self.tokenizer(prompt, return_tensors="pt")
+        outputs = self.model.generate(
+            **inputs, 
+            max_new_tokens=300,
+            num_return_sequences=1,
+            do_sample=True,
+            temperature=0.7,
+            top_p=0.9
+        )
+        
+        response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response_parts = response.split(prompt.split("\n\n")[-1])
+        return response_parts[-1].strip() if len(response_parts) > 1 else response.strip()
+
+    def translate_text(self, text, target_language):
+        if not text:
+            return "No text provided for translation."
+        
+        params = {'api-version': '3.0', 'to': target_language}
+        headers = {
+            'Ocp-Apim-Subscription-Key': AZURE_TRANSLATOR_KEY,
+            'Ocp-Apim-Subscription-Region': AZURE_TRANSLATOR_REGION,
+            'Content-Type': 'application/json'
+        }
+        body = [{'text': text}]
+        
+        try:
+            response = requests.post(AZURE_TRANSLATOR_ENDPOINT, params=params, headers=headers, json=body)
+            response_data = response.json()
+            return response_data[0]['translations'][0]['text']
+        except Exception as e:
+            return f"Error translating text: {str(e)}"
+
+# Create Gradio Interface
+def create_interface():
+    assistant = LegalEaseAssistant()
+
+    with gr.Blocks(title="LegalEase: AI Legal Assistant") as demo:
+        gr.Markdown("# 📜 LegalEase: AI-Powered Legal Document Assistant")
+        
+        with gr.Tabs():
+            for task_name, task_type in [("Simplify Language", "simplify"), ("Document Summary", "summary"), 
+                                         ("Key Terms", "key_terms"), ("Risk Analysis", "risk")]:
+                with gr.Tab(task_name):
+                    with gr.Row():
+                        input_file = gr.File(file_types=['txt', 'pdf', 'image'], label="Upload Document")
+                        input_text = gr.Textbox(label="Or Paste Text", lines=3)
+                    output_text = gr.Textbox(label="Generated Output", lines=6)
+                    
+                    language_dropdown = gr.Dropdown(choices=["en", "hi", "mr"], value="en", label="Translate To")
+                    translated_output = gr.Textbox(label="Translated Output", lines=6)
+                    btn = gr.Button(f"Generate {task_name}")
+
+                    def handler(file, text, language):
+                        input_source = file or text
+                        if not input_source:
+                            return "", ""
+                        generated_text = assistant.generate_response(input_source, task_type)
+                        translated_text = assistant.translate_text(generated_text, language) if language != "en" else generated_text
+                        return generated_text, translated_text
+                    
+                    btn.click(fn=handler, inputs=[input_file, input_text, language_dropdown], outputs=[output_text, translated_output])
+    
+    return demo
+
+# Create and launch the app
+demo = create_interface()
+
+if __name__ == "__main__":
+    demo.launch()