Update app.py

app.py

@@ -2,45 +2,75 @@ import os
 import torch
 from fastapi import FastAPI, Request
 from fastapi.responses import JSONResponse
-from fastapi.staticfiles import StaticFiles
 from fastapi.middleware.cors import CORSMiddleware
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from transformers import pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import time
+import warnings
+
+# Suppress specific warnings
+warnings.filterwarnings("ignore", category=FutureWarning, module="transformers.utils.hub")
+
+# Configure environment variables for cache
+os.environ["HF_HOME"] = os.getenv("HF_HOME", "/app/cache/huggingface")
+os.environ["MPLCONFIGDIR"] = os.getenv("MPLCONFIGDIR", "/app/cache/matplotlib")
 
 # Ensure cache directories exist
-os.makedirs(os.getenv('HUGGINGFACE_HUB_CACHE', '/app/cache/huggingface'), exist_ok=True)
-os.makedirs(os.getenv('MPLCONFIGDIR', '/app/cache/matplotlib'), exist_ok=True)
+os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+os.makedirs(os.environ["MPLCONFIGDIR"], exist_ok=True)
 
 # Initialize FastAPI app
 app = FastAPI()
 
-# Model loading function without autoawq
+def log_message(message: str):
+    """Helper function for logging"""
+    print(f"[{time.strftime('%Y-%m-%d %H:%M:%S')}] {message}")
+
 def load_model():
+    """Load the model with CPU optimization"""
     model_name = "trillionlabs/Trillion-7B-preview-AWQ"
     
-    # Load tokenizer with error handling
+    log_message("Loading tokenizer...")
     try:
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
             trust_remote_code=True
         )
     except Exception as e:
-        print(f"Error loading tokenizer: {e}")
-        # Fallback to a more basic tokenizer if needed
+        log_message(f"Tokenizer loading failed: {e}")
+        # Fallback to LlamaTokenizer if available
         from transformers import LlamaTokenizer
         tokenizer = LlamaTokenizer.from_pretrained(model_name)
     
-    # Load model with CPU configuration
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="cpu",
-        torch_dtype=torch.float32,
-        trust_remote_code=True
-    )
+    log_message("Loading model...")
+    try:
+        # Try loading with IPEX optimization if available
+        try:
+            import intel_extension_for_pytorch as ipex
+            use_ipex = True
+        except ImportError:
+            use_ipex = False
+            log_message("IPEX not available, using standard CPU version")
+
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float32,
+            trust_remote_code=True
+        )
+        
+        if use_ipex:
+            log_message("Applying IPEX optimization...")
+            model = ipex.optimize(model)
+        
+        # Explicitly move to CPU
+        model = model.to("cpu")
+        model.eval()
+        
+    except Exception as e:
+        log_message(f"Model loading failed: {e}")
+        raise
     
-    # Create text generation pipeline
+    log_message("Creating pipeline...")
     text_generator = pipeline(
         "text-generation",
         model=model,
@@ -52,21 +82,25 @@ def load_model():
 
 # Load model
 try:
+    log_message("Starting model loading process...")
     text_generator = load_model()
+    log_message("Model loaded successfully")
 except Exception as e:
-    print(f"Failed to load model: {e}")
-    # You might want to exit here or load a smaller model instead
+    log_message(f"Critical error loading model: {e}")
     raise
 
-# API endpoint for text generation
+# API endpoints
 @app.post("/api/generate")
-async def generate_text(request: Request):
+async def api_generate(request: Request):
+    """API endpoint for text generation"""
     try:
         data = await request.json()
-        prompt = data.get("prompt", "")
-        max_length = min(int(data.get("max_length", 100)), 500)  # Limit to 500 tokens
+        prompt = data.get("prompt", "").strip()
+        if not prompt:
+            return JSONResponse({"error": "Prompt cannot be empty"}, status_code=400)
+            
+        max_length = min(int(data.get("max_length", 100)), 300)  # Conservative limit
         
-        # Generate text with timing
         start_time = time.time()
         outputs = text_generator(
             prompt,
@@ -75,23 +109,39 @@ async def generate_text(request: Request):
             temperature=0.7,
             top_k=50,
             top_p=0.95,
-            pad_token_id=0  # Might be needed for some models
+            pad_token_id=tokenizer.eos_token_id if hasattr(text_generator, 'tokenizer') else 0
         )
         generation_time = time.time() - start_time
         
         return JSONResponse({
             "generated_text": outputs[0]["generated_text"],
-            "generation_time": round(generation_time, 2),
+            "time_seconds": round(generation_time, 2),
+            "tokens_generated": len(text_generator.tokenizer.tokenize(outputs[0]["generated_text"]) if hasattr(text_generator, 'tokenizer') else None,
             "model": "Trillion-7B-preview-AWQ",
             "device": "cpu"
         })
     except Exception as e:
+        log_message(f"API Error: {e}")
         return JSONResponse({"error": str(e)}, status_code=500)
 
-# Gradio interface
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "model_loaded": text_generator is not None,
+        "device": "cpu",
+        "cache_path": os.environ["HF_HOME"]
+    }
+
+# Gradio Interface
 def gradio_generate(prompt, max_length=100):
+    """Function for Gradio interface generation"""
     try:
-        max_length = min(int(max_length), 500)  # Limit to 500 tokens
+        max_length = min(int(max_length), 300)  # Same conservative limit as API
+        if not prompt.strip():
+            return "Please enter a prompt"
+            
         outputs = text_generator(
             prompt,
             max_length=max_length,
@@ -99,69 +149,72 @@ def gradio_generate(prompt, max_length=100):
             temperature=0.7,
             top_k=50,
             top_p=0.95,
-            pad_token_id=0
+            pad_token_id=tokenizer.eos_token_id if hasattr(text_generator, 'tokenizer') else 0
         )
         return outputs[0]["generated_text"]
     except Exception as e:
+        log_message(f"Gradio Error: {e}")
         return f"Error generating text: {str(e)}"
 
-with gr.Blocks() as gradio_app:
+with gr.Blocks(title="Trillion-7B CPU Demo", theme=gr.themes.Default()) as gradio_app:
     gr.Markdown("""
-    # Trillion-7B-preview-AWQ Demo (CPU)
-    *Running on CPU with 16GB RAM - responses may be slow*
+    # 🚀 Trillion-7B-preview-AWQ (CPU Version)
+    *Running on CPU with optimized settings - responses may be slower than GPU versions*
     """)
     
     with gr.Row():
-        input_prompt = gr.Textbox(
-            label="Input Prompt", 
-            lines=5,
-            placeholder="Enter your prompt here..."
-        )
-        output_text = gr.Textbox(
-            label="Generated Text", 
-            lines=5,
-            interactive=False
-        )
+        with gr.Column():
+            input_prompt = gr.Textbox(
+                label="Your Prompt",
+                placeholder="Enter text here...",
+                lines=5,
+                max_lines=10
+            )
+            with gr.Row():
+                max_length = gr.Slider(
+                    label="Max Length",
+                    minimum=20,
+                    maximum=300,
+                    value=100,
+                    step=10
+                )
+                generate_btn = gr.Button("Generate", variant="primary")
+        with gr.Column():
+            output_text = gr.Textbox(
+                label="Generated Text",
+                lines=10,
+                interactive=False
+            )
     
-    with gr.Row():
-        length_slider = gr.Slider(
-            minimum=50,
-            maximum=500,
-            value=100,
-            step=10,
-            label="Max Length"
-        )
-        generate_btn = gr.Button("Generate", variant="primary")
-    
-    # Additional examples
-    examples = gr.Examples(
+    # Examples
+    gr.Examples(
         examples=[
-            ["Explain quantum computing in simple terms."],
-            ["Write a short poem about artificial intelligence."],
-            ["How do I make a good cup of coffee?"]
+            ["Explain quantum computing in simple terms"],
+            ["Write a haiku about artificial intelligence"],
+            ["What are the main benefits of renewable energy?"],
+            ["Suggest three ideas for a science fiction story"]
         ],
-        inputs=input_prompt
+        inputs=input_prompt,
+        label="Example Prompts"
     )
     
     generate_btn.click(
         fn=gradio_generate,
-        inputs=[input_prompt, length_slider],
+        inputs=[input_prompt, max_length],
         outputs=output_text
     )
 
 # Mount Gradio app
 app = gr.mount_gradio_app(app, gradio_app, path="/")
 
-# CORS middleware
+# CORS configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
-    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Health check endpoint
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy", "model_loaded": text_generator is not None}
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)