Update download_model.py

download_model.py

@@ -1,16 +1,33 @@
-# download_model.py
+from fastapi import FastAPI
+from transformers import pipeline
 from huggingface_hub import snapshot_download
 from pathlib import Path
+import os
 
-def download_model():
-    mistral_models_path = Path.home().joinpath('mistral_models', '7B-v0.3')
-    mistral_models_path.mkdir(parents=True, exist_ok=True)
-    
+app = FastAPI()
+
+# Define the path where the model will be downloaded
+mistral_models_path = Path.home().joinpath('mistral_models', '7B-v0.3')
+mistral_models_path.mkdir(parents=True, exist_ok=True)
+
+# Download the model if not already present
+if not (mistral_models_path / "params.json").exists():
     snapshot_download(
-        repo_id="mistralai/Mistral-7B-v0.3",
-        allow_patterns=["params.json", "consolidated.safetensors", "tokenizer.model.v3"],
-        local_dir=mistral_models_path
+        repo_id="mistralai/Mistral-7B-v0.3", 
+        allow_patterns=["params.json", "consolidated.safetensors", "tokenizer.model.v3"], 
+        local_dir=mistral_models_path,
+        token=os.getenv('HUGGINGFACE_HUB_TOKEN')  # Use the environment variable for authentication
     )
 
-if __name__ == "__main__":
-    download_model()
+# Load the model
+pipe_mistral = pipeline("text2text-generation", model=str(mistral_models_path))
+
+@app.get("/mistral")
+def mistral_endpoint(input: str):
+    output = pipe_mistral(input)
+    return {"output": output[0]["generated_text"]}
+
+@app.get("/")
+def greet_json():
+    return {"Hello": "World!"}
+