Update app.py

app.py

@@ -2,13 +2,25 @@ from fastapi import FastAPI, Request
 from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import os
 
 app = FastAPI()
 
+# Create cache directory
+os.makedirs("./model_cache", exist_ok=True)
+
 # Load model and tokenizer once at startup
 model_name = "distilgpt2"  # change this to your own model
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+try:
+    # Try to load from local cache first
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="./model_cache", local_files_only=False)
+    model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="./model_cache", local_files_only=False)
+except OSError as e:
+    print(f"Error loading model: {e}")
+    print("Attempting to download model directly...")
+    # If that fails, try downloading explicitly
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="./model_cache")
+    model = AutoModelForCausalLM.from_pretrained(model_name, cache_dir="./model_cache")
 
 class PromptRequest(BaseModel):
     prompt: str
@@ -26,3 +38,7 @@ async def generate_text(req: PromptRequest):
     )
     generated = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return {"generated_text": generated}
+
+@app.get("/")
+async def root():
+    return {"status": "API is running", "model": model_name}