Update NER load fucntion

app.py

@@ -43,14 +43,25 @@ class SearchRequest(BaseModel):
     max_results: int = 5
 
 qa_pipeline = None
+ner_pipeline = None
 
 def load_ner_model():
     """Load Named Entity Recognition model"""
-    return pipeline(
-        "ner",
-        model="dbmdz/bert-large-cased-finetuned-conll03-english",
-        grouped_entities=True
-    )
+    global ner_pipeline
+    try:
+        # Check if GPU is available
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        logger.info(f"Using device: {device}")
+
+        ner_pipeline = pipeline(
+            "ner",
+            model="dbmdz/bert-large-cased-finetuned-conll03-english",
+            device=0 if device == "cuda" else -1,
+            grouped_entities=True
+        )
+        logger.info("NER model loaded successfully!")
+    except Exception as e:
+        logger.error(f"Error loading NER model: {e}")
 
 # Function to load the local language model
 def load_model():
@@ -114,7 +125,7 @@ def generate_response(prompt: str,search_context: str) -> str:
 async def startup_event():
     """Load model on startup"""
     load_model()
-    app.state.ner_pipeline = load_ner_model()
+    load_ner_model()
 
 @app.get("/")
 async def root():
@@ -140,7 +151,7 @@ async def chat(request: ChatRequest):
         if request.use_search:
 
             # Extract entities for focused search
-            entities = app.state.ner_pipeline(request.prompt)
+            entities = ner_pipeline(request.prompt)
             logger.info(f"Identified entities: {entities}")
 
             # Create search query from entities
@@ -149,11 +160,10 @@ async def chat(request: ChatRequest):
                 if ent["entity_group"] in ["PER", "ORG", "LOC", "MISC"]
             ]
             search_query = " ".join(search_terms) if search_terms else request.prompt
-            print(f"Search query: {search_query}")
+            logger.info(f"Search query: {search_query}")
 
             search_results = search_web(search_query)
-            
-
+        
             search_context = "\n".join([
                     f"- {result['title']}: {result['body'][:200]}..."
                     for result in search_results[:min(len(search_results), 5)]
@@ -162,7 +172,7 @@ async def chat(request: ChatRequest):
         else:
             search_context = request.prompt
 
-        print(f"Search context: {search_context}")
+        logger.info(f"Search context: {search_context}")
         if search_context:
         # Generate response
             response = generate_response(