Update app.py

app.py

@@ -2,14 +2,13 @@ import os
 import logging
 import sys
 import torch
+import tempfile
+from pathlib import Path
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel
-from typing import List, Dict, Optional, Any
-from datasets import load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from typing import List, Optional
 import uvicorn
-import time
 
 # Configure logging
 logging.basicConfig(
@@ -19,13 +18,15 @@ logging.basicConfig(
 )
 logger = logging.getLogger(__name__)
 
-app = FastAPI()
+app = FastAPI(title="Chat API", description="Simple chat API for Hugging Face Space")
 
-# Global variables
-model = None
-tokenizer = None
-generator = None
-dataset = None
+# Create a directory for caching in the current working directory
+cache_dir = Path("./model_cache")
+cache_dir.mkdir(exist_ok=True)
+os.environ["TRANSFORMERS_CACHE"] = str(cache_dir.absolute())
+os.environ["HF_HOME"] = str(cache_dir.absolute())
+
+logger.info(f"Using cache directory: {cache_dir.absolute()}")
 
 # Pydantic models for request/response
 class ChatTurn(BaseModel):
@@ -39,8 +40,23 @@ class ChatRequest(BaseModel):
 class ChatResponse(BaseModel):
     response: str
 
-# Use a much smaller model suitable for Hugging Face Spaces
-MODEL_ID = "distilgpt2"  # Using a very small model for testing
+# Global variables
+model = None
+tokenizer = None
+generator = None
+dataset = None
+
+# Load a small model or use a fallback if loading fails
+MODEL_ID = "distilgpt2"  # Small model for testing
+
+# Fallback responses for when the model isn't available
+FALLBACK_RESPONSES = [
+    "I apologize, but I'm currently having trouble processing your request.",
+    "Sorry, I'm experiencing technical difficulties at the moment.",
+    "I'm unable to generate a proper response right now. Please try again later.",
+    "My language model is temporarily unavailable. Please check back soon.",
+    "I would like to help, but I'm having some technical issues. Please try again shortly."
+]
 
 # Error handler
 @app.exception_handler(Exception)
@@ -51,58 +67,88 @@ async def generic_exception_handler(request: Request, exc: Exception):
         content={"detail": f"Internal server error: {str(exc)}"}
     )
 
-# Load model on startup
-@app.on_event("startup")
-async def startup_event():
-    global model, tokenizer, generator, dataset
+def try_load_model():
+    """Attempt to load the model and tokenizer with appropriate error handling"""
+    global model, tokenizer, generator
     
     try:
-        logger.info(f"Loading model: {MODEL_ID}")
-        start_time = time.time()
+        # Import here to handle import errors gracefully
+        from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
         
-        # Load the tokenizer
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        logger.info(f"Tokenizer loaded in {time.time() - start_time:.2f} seconds")
+        logger.info(f"Loading tokenizer for {MODEL_ID}")
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_ID,
+            cache_dir=cache_dir,
+            local_files_only=False
+        )
+        logger.info("Tokenizer loaded successfully")
         
-        # Load the model with optimizations
+        logger.info(f"Loading model {MODEL_ID}")
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_ID,
+            cache_dir=cache_dir,
+            local_files_only=False,
             torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-            low_cpu_mem_usage=True,
-            device_map="auto" if torch.cuda.is_available() else None
+            low_cpu_mem_usage=True
         )
-        logger.info(f"Model loaded in {time.time() - start_time:.2f} seconds")
+        logger.info("Model loaded successfully")
         
-        # Create a text generation pipeline
         device = 0 if torch.cuda.is_available() else -1
-        generator = pipeline("text-generation", model=model, tokenizer=tokenizer, device=device)
-        logger.info(f"Generator pipeline created in {time.time() - start_time:.2f} seconds")
-        
-        # Try to load dataset
-        try:
-            logger.info("Loading dataset: lahiruchamika27/tia")
-            dataset = load_dataset("lahiruchamika27/tia")
-            logger.info("Dataset loaded successfully")
-        except Exception as e:
-            logger.error(f"Error loading dataset: {str(e)}")
-            logger.info("Continuing without dataset")
-            
-        logger.info(f"Startup completed in {time.time() - start_time:.2f} seconds")
+        logger.info(f"Creating generator pipeline (device: {device})")
+        generator = pipeline(
+            "text-generation", 
+            model=model, 
+            tokenizer=tokenizer, 
+            device=device
+        )
+        logger.info("Generator pipeline created successfully")
+        return True
     except Exception as e:
-        logger.error(f"Error during startup: {str(e)}", exc_info=True)
-        logger.info("API will still be available but might not function correctly")
+        logger.error(f"Error loading model: {str(e)}", exc_info=True)
+        return False
 
+def try_load_dataset():
+    """Attempt to load the dataset with appropriate error handling"""
+    global dataset
+    
+    try:
+        from datasets import load_dataset
+        logger.info("Loading dataset: lahiruchamika27/tia")
+        dataset = load_dataset("lahiruchamika27/tia", cache_dir=cache_dir)
+        logger.info("Dataset loaded successfully")
+        return True
+    except Exception as e:
+        logger.error(f"Error loading dataset: {str(e)}", exc_info=True)
+        return False
+
+# Startup event
+@app.on_event("startup")
+async def startup_event():
+    logger.info("Starting application")
+    # Try to load model but don't fail if it doesn't work
+    model_loaded = try_load_model()
+    dataset_loaded = try_load_dataset()
+    logger.info(f"Startup complete. Model loaded: {model_loaded}, Dataset loaded: {dataset_loaded}")
+
+# Simple text-only route
+@app.get("/")
+async def root():
+    return {"message": "Chat API is running. Use /api/chat for chat functionality."}
+
+# Chat endpoint
 @app.post("/api/chat", response_model=ChatResponse)
 async def chat(request: ChatRequest):
     logger.info(f"Received chat request: {request.message[:50]}...")
     
-    # Check if model is loaded
+    # If the model isn't loaded, return a fallback response
     if generator is None:
-        logger.error("Text generator not initialized")
-        raise HTTPException(status_code=500, detail="Text generation pipeline not initialized")
+        import random
+        fallback = random.choice(FALLBACK_RESPONSES)
+        logger.warning("Using fallback response because model is not loaded")
+        return ChatResponse(response=fallback)
     
     try:
-        # Format conversation
+        # Format conversation history
         if request.history:
             full_prompt = ""
             for turn in request.history:
@@ -117,8 +163,7 @@ async def chat(request: ChatRequest):
         
         logger.info(f"Generated prompt: {full_prompt[:100]}...")
         
-        # Generate response
-        start_time = time.time()
+        # Generate text
         outputs = generator(
             full_prompt,
             max_new_tokens=100,
@@ -126,14 +171,13 @@ async def chat(request: ChatRequest):
             top_p=0.9,
             do_sample=True
         )
-        logger.info(f"Text generated in {time.time() - start_time:.2f} seconds")
         
         # Extract response
         generated_text = outputs[0]['generated_text']
-        # Extract only the assistant's response
+        # Extract just the assistant's response
         response_text = generated_text[len(full_prompt):].strip()
         
-        # If empty or just whitespace, return a fallback message
+        # Fallback if response is empty
         if not response_text or response_text.isspace():
             response_text = "I'm sorry, I'm having trouble generating a response right now."
             
@@ -141,8 +185,8 @@ async def chat(request: ChatRequest):
         return ChatResponse(response=response_text)
     
     except Exception as e:
-        logger.error(f"Error generating response: {str(e)}", exc_info=True)
-        raise HTTPException(status_code=500, detail=f"Error generating response: {str(e)}")
+        logger.error(f"Error in chat endpoint: {str(e)}", exc_info=True)
+        return ChatResponse(response="I'm sorry, I encountered an error while processing your request.")
 
 @app.get("/api/examples")
 async def get_examples(count: int = 5, split: str = "train"):
@@ -161,18 +205,25 @@ async def get_examples(count: int = 5, split: str = "train"):
 
 @app.get("/health")
 async def health_check():
-    system_info = {
+    return {
         "status": "ok",
         "model_loaded": model is not None,
         "tokenizer_loaded": tokenizer is not None,
         "generator_loaded": generator is not None,
         "dataset_loaded": dataset is not None,
-        "model_name": MODEL_ID,
-        "torch_device": "cuda" if torch.cuda.is_available() else "cpu",
-        "cuda_available": torch.cuda.is_available(),
-        "cuda_device_count": torch.cuda.device_count() if torch.cuda.is_available() else 0
+        "model_name": MODEL_ID if model is not None else None,
+        "device": "cuda" if torch.cuda.is_available() else "cpu",
+        "cache_dir": str(cache_dir)
+    }
+
+@app.get("/reload")
+async def reload_resources():
+    model_loaded = try_load_model()
+    dataset_loaded = try_load_dataset()
+    return {
+        "model_reloaded": model_loaded,
+        "dataset_reloaded": dataset_loaded
     }
-    return system_info
 
 if __name__ == "__main__":
     port = int(os.environ.get("PORT", 7860))