Update app.py

app.py

@@ -1,15 +1,22 @@
 from fastapi import FastAPI, HTTPException
 from transformers import AutoModelForCausalLM, AutoTokenizer
+from huggingface_hub import hf_hub_download
 import os
 
 app = FastAPI()
 
 cache_dir = os.path.expanduser("~/.cache/huggingface/")
 model_folder = os.path.join(cache_dir, "TheBloke/Mistral-7B-v0.1-GGUF")
+
+# Check if the model is present in the cache directory
 if not os.path.exists(model_folder):
-    raise ValueError("Model not found in cache directory. Please download the model.")
+    # Download the model if not present
+    try:
+        hf_hub_download(repo_id="TheBloke/Mistral-7B-v0.1-GGUF", filename="mistral-7b-v0.1.Q4_K_M.gguf")
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Failed to download model: {str(e)}")
 
-# Load the tokenizer and model asynchronously
+# Load the tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(model_folder)
 model = AutoModelForCausalLM.from_pretrained(model_folder)
 
@@ -22,4 +29,4 @@ async def generate_text():
         generated_texts = tokenizer.batch_decode(output, skip_special_tokens=True)
         return generated_texts
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))