barathm2001 commited on
Commit
98d251f
·
verified ·
1 Parent(s): 447816d

Upload 3 files

Browse files
Files changed (2) hide show
  1. app.py +15 -4
  2. requirements.txt +3 -5
app.py CHANGED
@@ -1,8 +1,8 @@
 
1
  import logging
2
  from fastapi import FastAPI, HTTPException
3
- from transformers import AutoModelForCausalLM, pipeline
4
  from peft import PeftModel, PeftConfig
5
- from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
6
 
7
  # Set up logging
8
  logging.basicConfig(level=logging.INFO)
@@ -21,17 +21,28 @@ async def load_model():
21
  global model, tokenizer, pipe
22
 
23
  try:
 
 
 
24
  logger.info("Loading PEFT configuration...")
25
  config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
26
 
27
  logger.info("Loading base model...")
28
- base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
 
 
 
 
29
 
30
  logger.info("Loading PEFT model...")
31
  model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
32
 
33
  logger.info("Loading tokenizer...")
34
- tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
 
 
 
 
35
 
36
  logger.info("Creating pipeline...")
37
  pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
1
+ import os
2
  import logging
3
  from fastapi import FastAPI, HTTPException
4
+ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
5
  from peft import PeftModel, PeftConfig
 
6
 
7
  # Set up logging
8
  logging.basicConfig(level=logging.INFO)
 
21
  global model, tokenizer, pipe
22
 
23
  try:
24
+ # Get Hugging Face token from environment variable
25
+ hf_token = os.environ.get("HUGGINGFACE_TOKEN")
26
+
27
  logger.info("Loading PEFT configuration...")
28
  config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
29
 
30
  logger.info("Loading base model...")
31
+ base_model = AutoModelForCausalLM.from_pretrained(
32
+ "mistralai/Mistral-7B-Instruct-v0.3",
33
+ token=hf_token if hf_token else None,
34
+ use_auth_token=True if not hf_token else None
35
+ )
36
 
37
  logger.info("Loading PEFT model...")
38
  model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
39
 
40
  logger.info("Loading tokenizer...")
41
+ tokenizer = AutoTokenizer.from_pretrained(
42
+ "mistralai/Mistral-7B-Instruct-v0.3",
43
+ token=hf_token if hf_token else None,
44
+ use_auth_token=True if not hf_token else None
45
+ )
46
 
47
  logger.info("Creating pipeline...")
48
  pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
requirements.txt CHANGED
@@ -1,11 +1,9 @@
1
  fastapi==0.103.0
2
- requests==2.27.*
3
  uvicorn[standard]==0.17.*
4
  torch>=1.13.0
5
- transformers>=4.34.0,<5.0.0
6
  numpy<2
7
- peft>=0.7.0
8
  accelerate>=0.24.1,<0.25.0
9
  huggingface_hub>=0.16.4,<0.18.0
10
- tokenizers>=0.14.0,<0.15.0
11
- git+https://github.com/mistralai/mistral-common.git@main
 
1
  fastapi==0.103.0
 
2
  uvicorn[standard]==0.17.*
3
  torch>=1.13.0
4
+ transformers>=4.34.0,<4.35.0
5
  numpy<2
6
+ peft>=0.6.0,<0.7.0
7
  accelerate>=0.24.1,<0.25.0
8
  huggingface_hub>=0.16.4,<0.18.0
9
+ tokenizers>=0.14.0,<0.15.0