barathm2001 commited on
Commit
93a3ab4
·
verified ·
1 Parent(s): cec48c7

Upload 3 files

Browse files
Files changed (2) hide show
  1. app.py +12 -2
  2. requirements.txt +1 -0
app.py CHANGED
@@ -1,17 +1,27 @@
 
1
  from fastapi import FastAPI
2
  from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
3
  from peft import PeftModel, PeftConfig
 
 
 
 
 
 
 
 
 
4
 
5
  # Initialize FastAPI app
6
  app = FastAPI()
7
 
8
  # Load PEFT model configuration and base model
9
  config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
10
- base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
11
  model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
12
 
13
  # Load tokenizer
14
- tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
15
 
16
  # Create the pipeline
17
  pipe = pipeline("text2sql", model=model, tokenizer=tokenizer)
 
1
+ import os
2
  from fastapi import FastAPI
3
  from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer
4
  from peft import PeftModel, PeftConfig
5
+ from huggingface_hub import login
6
+
7
+ # Get the Hugging Face token from the environment variable
8
+ huggingface_token = os.getenv("HUGGING_FACE_TOKEN")
9
+ if huggingface_token is None:
10
+ raise ValueError("HUGGING_FACE_TOKEN environment variable is not set")
11
+
12
+ # Login to Hugging Face Hub
13
+ login(token=huggingface_token)
14
 
15
  # Initialize FastAPI app
16
  app = FastAPI()
17
 
18
  # Load PEFT model configuration and base model
19
  config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
20
+ base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", use_auth_token=True)
21
  model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
22
 
23
  # Load tokenizer
24
+ tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3", use_auth_token=True)
25
 
26
  # Create the pipeline
27
  pipe = pipeline("text2sql", model=model, tokenizer=tokenizer)
requirements.txt CHANGED
@@ -6,3 +6,4 @@ torch>=1.13.0
6
  transformers==4.*
7
  numpy<2
8
  peft==0.11.1
 
 
6
  transformers==4.*
7
  numpy<2
8
  peft==0.11.1
9
+ huggingface-hub==0.15.1