Update app.py

app.py

@@ -1,172 +1,278 @@
 #!/usr/bin/env python3
 """
-Indian Legal AI Assistant using Hugging Face LLM
-This app provides a chat interface for querying Indian laws using the
-invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF model.
+Indian Legal AI Assistant using a GGUF model with llama-cpp-python.
+
+Model:
+invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF
 """
 
 import os
+import traceback
+
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
 
+
+# -----------------------------
 # Model configuration
+# -----------------------------
+
 MODEL_REPO = "invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF"
 MODEL_FILE = "ambuj-tripathi-indian-legal-llama.Q4_K_M.gguf"
 
+# Hugging Face Spaces usually provides limited CPU/RAM on free hardware.
+# These defaults are conservative.
+N_CTX = int(os.getenv("N_CTX", "2048"))
+N_THREADS = int(os.getenv("N_THREADS", "4"))
+N_GPU_LAYERS = int(os.getenv("N_GPU_LAYERS", "0"))
+
 # Global model instance
 llm = None
 
 
+# -----------------------------
+# Model loading
+# -----------------------------
+
 def load_model():
-    """Load the GGUF model from Hugging Face Hub"""
+    """Download and load the GGUF model once."""
     global llm
-    if llm is None:
-        print(f"Downloading model from {MODEL_REPO}...")
-        model_path = hf_hub_download(
-            repo_id=MODEL_REPO,
-            filename=MODEL_FILE,
-            cache_dir="./models"
-        )
-        print(f"Model downloaded to: {model_path}")
-        print("Loading model into memory...")
-        llm = Llama(
-            model_path=model_path,
-            n_ctx=2048,  # Context window
-            n_threads=4,  # CPU threads
-            n_gpu_layers=0,  # Set to 0 for CPU-only, increase if GPU available
-        )
-        print("Model loaded successfully!")
+
+    if llm is not None:
+        return llm
+
+    print(f"Downloading model from {MODEL_REPO}...")
+    model_path = hf_hub_download(
+        repo_id=MODEL_REPO,
+        filename=MODEL_FILE,
+        cache_dir="./models",
+    )
+
+    print(f"Model downloaded to: {model_path}")
+    print("Loading model into memory...")
+
+    llm = Llama(
+        model_path=model_path,
+        n_ctx=N_CTX,
+        n_threads=N_THREADS,
+        n_gpu_layers=N_GPU_LAYERS,
+        verbose=False,
+    )
+
+    print("Model loaded successfully.")
     return llm
 
 
-def chat_with_legal_ai(message, history):
-    """
-    Process user message and generate response using the Indian Legal LLM
-    
-    Args:
-        message: User's input message
-        history: Chat history (list of [user_msg, bot_msg] pairs)
-    
-    Returns:
-        Response from the LLM
-    """
-    try:
-        model = load_model()
-        
-        # Build conversation context
-        conversation = ""
-        if history:
-            for user_msg, bot_msg in history:
-                conversation += f"User: {user_msg}\nAssistant: {bot_msg}\n"
-        
-        # Add current message
-        prompt = f"""You are an expert Indian legal AI assistant. You have deep knowledge of Indian laws, acts, and legal procedures. Provide accurate, helpful, and concise answers to legal queries.
+# -----------------------------
+# Prompting / inference
+# -----------------------------
+
+def build_prompt(message, history):
+    """Build a simple instruction prompt from chat history."""
+
+    conversation = ""
+
+    if history:
+        for user_msg, bot_msg in history:
+            conversation += f"User: {user_msg}\nAssistant: {bot_msg}\n"
+
+    prompt = f"""You are an expert Indian legal AI assistant.
+
+You have knowledge of Indian laws, acts, legal procedures, and general legal concepts.
+Provide accurate, helpful, and concise information.
+
+Important:
+- Do not claim to be a lawyer.
+- Do not present your answer as formal legal advice.
+- Encourage the user to consult a qualified lawyer for specific legal matters.
 
 {conversation}User: {message}
 Assistant:"""
-        
-        # Generate response
+
+    return prompt
+
+
+def generate_legal_response(message, history):
+    """Generate a response from the local GGUF model."""
+
+    if not message or not message.strip():
+        return "Please enter a legal question."
+
+    try:
+        model = load_model()
+        prompt = build_prompt(message.strip(), history)
+
         response = model(
             prompt,
             max_tokens=512,
             temperature=0.7,
             top_p=0.95,
             echo=False,
-            stop=["User:", "\n\n\n"]
+            stop=["User:", "\nUser:", "\n\n\n"],
         )
-        
-        return response['choices'][0]['text'].strip()
-    
-    except Exception as e:
-        return f"Error: {str(e)}\n\nPlease ensure the model is properly downloaded and configured."
 
+        text = response["choices"][0]["text"].strip()
+
+        if not text:
+            return "I could not generate a response. Please try rephrasing your question."
+
+        return text
+
+    except Exception as error:
+        print("Error during generation:")
+        traceback.print_exc()
+
+        return (
+            "The app encountered an error while loading or running the model.\n\n"
+            f"Error details: {str(error)}"
+        )
+
+
+# -----------------------------
+# Gradio handlers
+# -----------------------------
+
+def respond(message, history):
+    """
+    Gradio chat handler.
+
+    Input:
+    - message: latest user message
+    - history: list of (user_message, assistant_message)
+
+    Output:
+    - empty textbox
+    - updated chatbot history
+    """
+
+    if history is None:
+        history = []
+
+    bot_reply = generate_legal_response(message, history)
+    history.append((message, bot_reply))
+
+    return "", history
+
+
+def clear_chat():
+    """Clear the chat history."""
+    return []
+
+
+# -----------------------------
+# Gradio UI
+# -----------------------------
 
 def create_gradio_interface():
-    """Create and return the Gradio chat interface"""
-    
-    # Custom CSS for better appearance
+    """Create and return the Gradio chat interface."""
+
     custom_css = """
     .container {
         max-width: 900px;
         margin: auto;
         padding: 20px;
     }
+
     #title {
         text-align: center;
         color: #1e3a8a;
     }
     """
-    
+
     with gr.Blocks(css=custom_css, title="Indian Legal AI Assistant") as demo:
         gr.Markdown(
             """
             # 🏛️ Indian Legal AI Assistant
-            
-            Ask questions about Indian laws, acts, and legal procedures. This AI assistant is powered by 
-            the **Ambuj Tripathi Indian Legal Llama** model, trained on Indian legal documents.
-            
+
+            Ask questions about Indian laws, acts, and legal procedures.
+
+            This AI assistant uses the **Ambuj Tripathi Indian Legal Llama GGUF** model.
+
             **Examples:**
             - What is the Indian Penal Code?
-            - Explain Section 377 of IPC
-            - What are the grounds for divorce under Hindu Marriage Act?
-            - Explain the Right to Information Act
+            - Explain Section 377 of IPC.
+            - What are the grounds for divorce under the Hindu Marriage Act?
+            - Explain the Right to Information Act.
             """,
-            elem_id="title"
+            elem_id="title",
         )
-        
+
         chatbot = gr.Chatbot(
             height=500,
             label="Chat History",
             show_label=True,
-            elem_id="chatbot"
+            elem_id="chatbot",
+            type="tuples",
         )
-        
+
         with gr.Row():
             msg = gr.Textbox(
                 label="Your Question",
                 placeholder="Ask about Indian laws...",
                 lines=2,
-                scale=4
+                scale=4,
+            )
+
+            submit = gr.Button(
+                "Send",
+                variant="primary",
+                scale=1,
             )
-            submit = gr.Button("Send", variant="primary", scale=1)
-        
+
         with gr.Row():
             clear = gr.Button("Clear Chat")
-        
+
         gr.Markdown(
             """
             ---
-            **Disclaimer:** This AI assistant provides general information only and should not be considered 
-            legal advice. For specific legal matters, please consult a qualified legal professional.
-            
+
+            **Disclaimer:** This assistant provides general legal information only.
+            It is not a substitute for advice from a qualified legal professional.
+            For specific legal matters, please consult a lawyer.
+
             **Model:** [invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF](https://huggingface.co/invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF)
             """
         )
-        
-        # Set up event handlers
-        msg.submit(chat_with_legal_ai, [msg, chatbot], [chatbot])
-        submit.click(chat_with_legal_ai, [msg, chatbot], [chatbot])
-        msg.submit(lambda: "", None, [msg])
-        submit.click(lambda: "", None, [msg])
-        clear.click(lambda: None, None, [chatbot])
-    
+
+        msg.submit(
+            respond,
+            inputs=[msg, chatbot],
+            outputs=[msg, chatbot],
+        )
+
+        submit.click(
+            respond,
+            inputs=[msg, chatbot],
+            outputs=[msg, chatbot],
+        )
+
+        clear.click(
+            clear_chat,
+            inputs=None,
+            outputs=chatbot,
+        )
+
     return demo
 
 
+# -----------------------------
+# App entry point
+# -----------------------------
+
 if __name__ == "__main__":
     print("Starting Indian Legal AI Assistant...")
     print(f"Using model: {MODEL_REPO}/{MODEL_FILE}")
-    
-    # Pre-load the model to avoid delays on first query
-    print("\nPre-loading model (this may take a few minutes on first run)...")
-    load_model()
-    
-    # Launch the Gradio interface
+
+    # Do not preload the model here.
+    # On Hugging Face Spaces, preloading can cause startup timeout or memory issues.
+    # The model will load on the first user message instead.
+
     demo = create_gradio_interface()
+
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
         share=False,
-        show_error=True
+        show_error=True,
     )