Upload 3 files

app.py

@@ -1,7 +1,9 @@
 import logging
+import os
 from fastapi import FastAPI, HTTPException
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from peft import PeftModel, PeftConfig
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -15,22 +17,27 @@ model = None
 tokenizer = None
 pipe = None
 
+# Get the Hugging Face token from environment variable
+hf_token = os.environ.get("HUGGINGFACE_TOKEN")
+if not hf_token:
+    raise ValueError("HUGGINGFACE_TOKEN environment variable is not set")
+
 @app.on_event("startup")
 async def load_model():
     global model, tokenizer, pipe
     
     try:
         logger.info("Loading PEFT configuration...")
-        config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
+        config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval", token=hf_token)
         
         logger.info("Loading base model...")
-        base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+        base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", token=hf_token)
         
         logger.info("Loading PEFT model...")
-        model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
+        model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval", token=hf_token)
 
         logger.info("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+        tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", token=hf_token)
 
         logger.info("Creating pipeline...")
         pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)

requirements.txt

@@ -2,9 +2,9 @@ fastapi==0.103.0
 requests==2.27.*
 uvicorn[standard]==0.17.*
 torch>=1.13.0
-transformers>=4.34.0,<5.0.0
+transformers>=4.36.0,<5.0.0
 numpy<2
-peft>=0.7.0
+peft>=0.8.0
 accelerate>=0.24.1,<0.25.0
 huggingface_hub>=0.16.4,<0.18.0
 tokenizers>=0.14.0,<0.15.0