Updated package v4.0

Author: UtkarshTheDev

locallab/__init__.py

@@ -2,7 +2,7 @@
 LocalLab: Run LLMs locally with a friendly API similar to OpenAI
 """
 
-__version__ = "0.3.9"
+__version__ = "0.4.0"
 
 from typing import Dict, Any, Optional
 

locallab/model_manager.py

@@ -176,6 +176,19 @@ def _apply_optimizations(self, model: AutoModelForCausalLM) -> AutoModelForCausa
                     logger.warning(
                         f"BetterTransformer optimization failed: {str(e)}")
 
+            # Only apply Flash Attention if explicitly enabled and not empty
+            if os.environ.get('LOCALLAB_ENABLE_FLASH_ATTENTION', '').lower() not in ('false', '0', 'none', ''):
+                try:
+                    from flash_attn import FlashAttention
+                    model = FlashAttention(model)
+                    logger.info("Flash Attention optimization applied")
+                except ImportError:
+                    logger.warning(
+                        "Flash Attention not available - install 'flash-attn' for this feature")
+                except Exception as e:
+                    logger.warning(
+                        f"Flash Attention optimization failed: {str(e)}")
+
             return model
         except Exception as e:
             logger.warning(f"Some optimizations could not be applied: {str(e)}")
@@ -219,10 +232,13 @@ async def load_model(self, model_id: str) -> bool:
                     **config
                 )
 
-                # Move model to the appropriate device
+                # Move model only if quantization is disabled
                 if not ENABLE_QUANTIZATION or str(ENABLE_QUANTIZATION).lower() in ('false', '0', 'none', ''):
                     device = "cuda" if torch.cuda.is_available() else "cpu"
+                    logger.info(f"Moving model to {device}")
                     self.model = self.model.to(device)
+                else:
+                    logger.info("Skipping device move for quantized model - using device_map='auto'")
 
                 # Capture model parameters after loading
                 model_architecture = self.model.config.architectures[0] if hasattr(self.model.config, 'architectures') else 'Unknown'

locallab/ui/banners.py

@@ -187,7 +187,7 @@ def print_api_docs():
     }}
 
   • Example:
-    curl -X POST "http://localhost:8000/generate" \\
+    curl -X POST "<server-ngrok-public-url>/generate" \\
     -H "Content-Type: application/json" \\
     -d '{{"prompt": "Write a story about a dragon", "max_tokens": 100}}'
 
@@ -204,35 +204,35 @@ def print_api_docs():
     }}
 
   • Example:
-    curl -X POST "http://localhost:8000/chat" \\
+    curl -X POST "<server-ngrok-public-url>/chat" \\
     -H "Content-Type: application/json" \\
     -d '{{"messages": [{{"role": "user", "content": "Hello, who are you?"}}]}}'
 
 📦 Model Management Endpoints:
 
 1️⃣ /models - List available models
   • GET
-  • Example: curl "http://localhost:8000/models"
+  • Example: curl "<server-ngrok-public-url>/models"
 
 2️⃣ /models/load - Load a specific model
   • POST with JSON body: {{ "model_id": "microsoft/phi-2" }}
   • Example:
-    curl -X POST "http://localhost:8000/models/load" \\
+    curl -X POST "<server-ngrok-public-url>/models/load" \\
     -H "Content-Type: application/json" \\
     -d '{{"model_id": "microsoft/phi-2"}}'
 
 ℹ️ System Endpoints:
 
 1️⃣ /system/info - Get system information
   • GET
-  • Example: curl "http://localhost:8000/system/info"
+  • Example: curl "<server-ngrok-public-url>/system/info"
 
 2️⃣ /system/resources - Get detailed system resources
   • GET
-  • Example: curl "http://localhost:8000/system/resources"
+  • Example: curl "<server-ngrok-public-url>/system/resources"
 
 3️⃣ /docs - Interactive API documentation (Swagger UI)
-  • Open in browser: http://localhost:8000/docs
+  • Open in browser: <server-ngrok-public-url>/docs
 
 {Fore.CYAN}════════════════════════════════════════════════════════════════════════{Style.RESET_ALL}
 """

setup.py

@@ -5,7 +5,7 @@
 
 setup(
     name="locallab",
-    version="0.3.9",
+    version="0.4.0",
     packages=find_packages(include=["locallab", "locallab.*"]),
     install_requires=[
         "fastapi>=0.95.0,<1.0.0",