Upload 3 files

app.py

@@ -1,17 +1,27 @@
+import os
 from fastapi import FastAPI
 from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
 from peft import PeftModel, PeftConfig
+from huggingface_hub import login
+
+# Get the Hugging Face token from the environment variable
+huggingface_token = os.getenv("HUGGING_FACE_TOKEN")
+if huggingface_token is None:
+    raise ValueError("HUGGING_FACE_TOKEN environment variable is not set")
+
+# Login to Hugging Face Hub
+login(token=huggingface_token)
 
 # Initialize FastAPI app
 app = FastAPI()
 
 # Load PEFT model configuration and base model
 config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
-base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", use_auth_token=True)
 model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
 
 # Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", use_auth_token=True)
 
 # Create the pipeline
 pipe = pipeline("text2sql", model=model, tokenizer=tokenizer)

requirements.txt

@@ -6,3 +6,4 @@ torch>=1.13.0
 transformers==4.*
 numpy<2
 peft==0.11.1
+huggingface-hub==0.15.1