Update app.py

app.py

@@ -6,52 +6,106 @@ import faiss
 import numpy as np
 import json
 import torch
+import requests
+import threading
+from queue import Queue
+import logging
+import pynvml
 
-# Load mock legal corpus
+# Configure logging
+logging.basicConfig(level=logging.DEBUG)
+logger = logging.getLogger(__name__)
+
+# Initialize NVML for GPU debugging
+try:
+    pynvml.nvmlInit()
+    device_count = pynvml.nvmlDeviceGetCount()
+    logger.info(f"NVML Initialized. GPU Count: {device_count}")
+    for i in range(device_count):
+        handle = pynvml.nvmlDeviceGetHandleByIndex(i)
+        name = pynvml.nvmlDeviceGetName(handle)
+        logger.info(f"GPU {i}: {name}")
+except pynvml.NVMLError as e:
+    logger.error(f"NVML Initialization Failed: {str(e)}")
+    raise RuntimeError("Cannot initialize NVML. Check NVIDIA drivers.")
+
+# Verify CUDA
+if not torch.cuda.is_available():
+    logger.error("CUDA not available")
+    raise RuntimeError("No GPU detected. Ensure H200 GPU is available.")
+logger.info(f"CUDA Version: {torch.version.cuda}")
+logger.info(f"GPU Detected: {torch.cuda.get_device_name(0)}")
+logger.info(f"Device Count: {torch.cuda.device_count()}")
+
+# Load legal corpus
 with open("legal_corpus.json", "r", encoding="utf-8") as f:
     corpus = json.load(f)
 documents = [item["text"] for item in corpus]
 
-# Initialize sentence transformer for embeddings
-embedder = SentenceTransformer("all-MiniLM-L6-v2")  # Lightweight embedder
+# Initialize sentence transformer (GPU)
+embedder = SentenceTransformer("all-MiniLM-L6-v2", device="cuda")
 embeddings = embedder.encode(documents, convert_to_numpy=True)
 
-# Initialize FAISS index
+# Initialize FAISS-GPU
 dimension = embeddings.shape[1]
 index = faiss.IndexFlatL2(dimension)
 index.add(embeddings)
 
-# Initialize GLiNER model
+# Initialize GLiNER (GPU)
 gliner_model = GLiNER.from_pretrained("NAMAA-Space/gliner_arabic-v2.1", load_tokenizer=True)
+gliner_model = gliner_model.cuda()
+
+# Initialize LLM (default to Qwen2-7B-Instruct-AWQ)
+use_qwq_32b = False  # Set to True if H200 detection is fixed
+model_name = "Qwen/Qwen2-7B-Instruct-AWQ" if not use_qwq_32b else "Qwen/QwQ-32B"
+try:
+    llm = LLM(
+        model=model_name,
+        quantization="awq",
+        max_model_len=4096,
+        gpu_memory_utilization=0.9,
+        device="cuda"
+    )
+    logger.info(f"Loaded LLM: {model_name}")
+except Exception as e:
+    logger.error(f"Failed to initialize LLM: {str(e)}")
+    raise
 
-# Initialize QwQ-32B
-llm = LLM(
-    model="Qwen/QwQ-32B",
-    quantization="awq",
-    max_model_len=4096,
-    gpu_memory_utilization=0.9
-)
 sampling_params = SamplingParams(temperature=0.7, max_tokens=512)
 
-def retrieve_documents(query, k=2):
-    """Retrieve top-k relevant documents using FAISS."""
-    query_embedding = embedder.encode([query], convert_to_numpy=True)
-    distances, indices = index.search(query_embedding, k)
-    return [documents[idx] for idx in indices[0]]
+def fetch_external_legal_data(query, queue):
+    """Fetch external legal data via HTTP request (mock API)."""
+    try:
+        response = requests.get(
+            "https://api.example.com/legal",
+            params={"query": query},
+            timeout=5
+        )
+        response.raise_for_status()
+        queue.put(response.json().get("text", "No external data found"))
+    except requests.RequestException:
+        queue.put("Failed to fetch external data")
 
-def run_ner(text, entity_types):
+def run_ner(text, entity_types, queue):
     """Run NER with gliner_arabic-v2.1."""
     if not text or not entity_types:
-        return []
+        queue.put([])
+        return
     entity_list = [e.strip() for e in entity_types.split(",")]
     entities = gliner_model.predict_entities(text, entity_list, threshold=0.5)
-    return [{"text": e["text"], "label": e["label"], "score": round(e["score"], 2)} for e in entities]
+    queue.put([{"text": e["text"], "label": e["label"], "score": round(e["score"], 2)} for e in entities])
 
-def generate_legal_insight(text, entities, retrieved_docs):
-    """Generate insight with QwQ-32B using RAG."""
+def retrieve_documents(query, k=2):
+    """Retrieve top-k documents using FAISS-GPU."""
+    query_embedding = embedder.encode([query], convert_to_numpy=True)
+    distances, indices = index.search(query_embedding, k)
+    return [documents[idx] for idx in indices[0]]
+
+def generate_legal_insight(text, entities, retrieved_docs, external_data):
+    """Generate insight with LLM using RAG."""
     entity_str = ", ".join([f"{e['text']} ({e['label']})" for e in entities])
-    context = "\n".join(retrieved_docs)
-    prompt = f"""You are a legal assistant for Arabic law. Using the following context and extracted entities, provide a concise legal insight (e.g., summary or explanation). Ensure the response is grounded in the context and entities.
+    context = "\n".join(retrieved_docs) + "\nExternal Data: " + external_data
+    prompt = f"""You are a legal assistant for Arabic law. Using the following context, extracted entities, and external data, provide a concise legal insight.
 
 Context:
 {context}
@@ -67,21 +121,31 @@ Insight:"""
     return outputs[0].outputs[0].text
 
 def main_interface(text, entity_types):
-    """Main Gradio interface."""
-    # Run NER
-    ner_result = run_ner(text, entity_types)
+    """Main Gradio interface with threading."""
+    ner_queue = Queue()
+    external_queue = Queue()
+
+    ner_thread = threading.Thread(target=run_ner, args=(text, entity_types, ner_queue))
+    external_thread = threading.Thread(target=fetch_external_legal_data, args=(text, external_queue))
+    
+    ner_thread.start()
+    external_thread.start()
     
-    # Retrieve relevant documents
+    ner_thread.join()
+    external_thread.join()
+
+    ner_result = ner_queue.get()
+    external_data = external_queue.get()
+
     retrieved_docs = retrieve_documents(text)
     
-    # Generate legal insight
-    insight = generate_legal_insight(text, ner_result, retrieved_docs)
+    insight = generate_legal_insight(text, ner_result, retrieved_docs, external_data)
     
-    return ner_result, retrieved_docs, insight
+    return ner_result, retrieved_docs, external_data, insight
 
 # Gradio interface
 with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
-    gr.Markdown("# Arabic Legal Demo: NER & RAG with GLiNER and QwQ-32B")
+    gr.Markdown("# Arabic Legal Demo: NER & RAG with GLiNER and LLM")
     with gr.Row():
         text_input = gr.Textbox(label="Arabic Legal Text", lines=5, placeholder="Enter Arabic legal text...")
         entity_types = gr.Textbox(
@@ -92,12 +156,13 @@ with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
     submit_btn = gr.Button("Analyze")
     ner_output = gr.JSON(label="Extracted Entities")
     docs_output = gr.Textbox(label="Retrieved Legal Context")
+    external_output = gr.Textbox(label="External Legal Data")
     insight_output = gr.Textbox(label="Legal Insight")
     
     submit_btn.click(
         fn=main_interface,
         inputs=[text_input, entity_types],
-        outputs=[ner_output, docs_output, insight_output]
+        outputs=[ner_output, docs_output, external_output, insight_output]
     )
 
 demo.launch()