Upload 4 files

app.py

@@ -1,8 +1,9 @@
-import os
 import logging
 from fastapi import FastAPI, HTTPException
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, pipeline
 from peft import PeftModel, PeftConfig
+from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
+from mistral_common.client import MistralChain
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -15,39 +16,35 @@ app = FastAPI()
 model = None
 tokenizer = None
 pipe = None
+mistral_chain = None
 
 @app.on_event("startup")
 async def load_model():
-    global model, tokenizer, pipe
+    global model, tokenizer, pipe, mistral_chain
     
     try:
-        # Get Hugging Face token from environment variable
-        hf_token = os.environ.get("HUGGINGFACE_TOKEN")
-        
         logger.info("Loading PEFT configuration...")
         config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
         
         logger.info("Loading base model...")
-        base_model = AutoModelForCausalLM.from_pretrained(
-            "mistralai/Mistral-7B-Instruct-v0.3", 
-            token=hf_token if hf_token else None,
-            use_auth_token=True if not hf_token else None
-        )
+        base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
         
         logger.info("Loading PEFT model...")
         model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
 
         logger.info("Loading tokenizer...")
-        tokenizer = AutoTokenizer.from_pretrained(
-            "mistralai/Mistral-7B-Instruct-v0.3", 
-            token=hf_token if hf_token else None,
-            use_auth_token=True if not hf_token else None
-        )
+        tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
+
+        logger.info("Creating MistralChain...")
+        mistral_chain = MistralChain(model, tokenizer)
 
         logger.info("Creating pipeline...")
         pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
         
         logger.info("Model, tokenizer, and pipeline loaded successfully.")
+    except ImportError as e:
+        logger.error(f"Error importing required modules. Please check your installation: {e}")
+        raise
     except Exception as e:
         logger.error(f"Error loading model or creating pipeline: {e}")
         raise
@@ -58,12 +55,12 @@ def home():
 
 @app.get("/generate")
 async def generate(text: str):
-    if not pipe:
+    if not mistral_chain:
         raise HTTPException(status_code=503, detail="Model not loaded")
     
     try:
-        output = pipe(text, max_length=100, num_return_sequences=1)
-        return {"output": output[0]['generated_text']}
+        output = mistral_chain.generate(text, max_tokens=100)
+        return {"output": output}
     except Exception as e:
         logger.error(f"Error during text generation: {e}")
         raise HTTPException(status_code=500, detail=f"Error during text generation: {str(e)}")

requirements.txt

@@ -1,9 +1,11 @@
 fastapi==0.103.0
+requests==2.27.*
 uvicorn[standard]==0.17.*
 torch>=1.13.0
-transformers>=4.34.0,<4.35.0
+transformers>=4.36.0,<5.0.0
 numpy<2
-peft>=0.6.0,<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
+tokenizers>=0.14.0,<0.15.0
+git+https://github.com/mistralai/mistral-common.git@main