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

Upload 3 files

Browse files
Files changed (1) hide show
  1. app.py +30 -12
app.py CHANGED
@@ -1,22 +1,37 @@
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,5 +39,8 @@ def home():
24
 
25
  @app.get("/generate")
26
  def generate(text: str):
27
- output = pipe(text)
28
- return {"output": output[0]['generated_text']}
 
 
 
 
1
  from fastapi import FastAPI
2
  from transformers import AutoModelForCausalLM
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.tokens.tokenizers.mistral import MistralTokenizer
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 = MistralTokenizer.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
 
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)}