barathm2001 commited on
Commit
21b7e63
·
verified ·
1 Parent(s): 7287f22

Upload 3 files

Browse files
Files changed (1) hide show
  1. app.py +13 -31
app.py CHANGED
@@ -1,37 +1,22 @@
1
  from fastapi import FastAPI
2
- from transformers import AutoModelForCausalLM, AutoTokenizer
3
-
4
- # Wrap problematic imports in try-except blocks
5
- try:
6
- from peft import PeftModel, PeftConfig
7
- except ImportError as e:
8
- print(f"Error importing from peft: {e}")
9
- raise
10
-
11
- try:
12
- from mistral_common.tokenizer import AutoMistralTokenizer
13
- except ImportError as e:
14
- print(f"Error importing from mistral_common: {e}")
15
- raise
16
 
17
  # Initialize FastAPI app
18
  app = FastAPI()
19
 
20
  # Load PEFT model configuration and base model
21
- try:
22
- config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
23
- base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
24
- model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
25
 
26
- # Load recommended tokenizer
27
- tokenizer = AutoMistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
28
 
29
- # Create the pipeline
30
- from transformers import pipeline
31
- pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
32
- except Exception as e:
33
- print(f"Error loading model or creating pipeline: {e}")
34
- raise
35
 
36
  @app.get("/")
37
  def home():
@@ -39,8 +24,5 @@ def home():
39
 
40
  @app.get("/generate")
41
  def generate(text: str):
42
- try:
43
- output = pipe(text)
44
- return {"output": output[0]['generated_text']}
45
- except Exception as e:
46
- return {"error": str(e)}
 
1
  from fastapi import FastAPI
2
+ from transformers import AutoModelForCausalLM
3
+ from mistral_common import MistralTokenizer # Hypothetical package, adjust based on actual package name and usage
4
+ from peft import PeftModel, PeftConfig
 
 
 
 
 
 
 
 
 
 
 
5
 
6
  # Initialize FastAPI app
7
  app = FastAPI()
8
 
9
  # Load PEFT model configuration and base model
10
+ config = PeftConfig.from_pretrained("frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
11
+ base_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
12
+ model = PeftModel.from_pretrained(base_model, "frankmorales2020/Mistral-7B-text-to-sql-flash-attention-2-dataeval")
 
13
 
14
+ # Load recommended tokenizer
15
+ tokenizer = MistralTokenizer.from_pretrained("mistralai/Mistral-7B-Instruct-v0.3")
16
 
17
+ # Create the pipeline
18
+ from transformers import pipeline
19
+ pipe = pipeline("text2sql", model=model, tokenizer=tokenizer)
 
 
 
20
 
21
  @app.get("/")
22
  def home():
 
24
 
25
  @app.get("/generate")
26
  def generate(text: str):
27
+ output = pipe(text)
28
+ return {"output": output[0]['generated_text']}