dohael commited on
Commit
d95047f
·
verified ·
1 Parent(s): a985ba5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +17 -12
app.py CHANGED
@@ -6,17 +6,21 @@ import torch
6
 
7
  app = FastAPI(title="Assistant IA Education Marocaine")
8
 
9
- # Chargement du modèle
10
  MODEL_NAME = "unsloth/mistral-7b-v0.3"
11
  LORA_NAME = "dohael/mistral-7b-education-maroc"
12
 
13
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
14
  model = AutoModelForCausalLM.from_pretrained(
15
  MODEL_NAME,
16
- torch_dtype=torch.float16,
17
- device_map="auto"
18
  )
19
- model = PeftModel.from_pretrained(model, LORA_NAME)
 
 
 
 
 
20
 
21
  class Question(BaseModel):
22
  question: str
@@ -28,14 +32,15 @@ def root():
28
  @app.post("/ask")
29
  def ask(q: Question):
30
  prompt = f"<s>[INST] {q.question} [/INST]"
31
- inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
32
- outputs = model.generate(
33
- **inputs,
34
- max_new_tokens=512,
35
- temperature=0.7,
36
- do_sample=True,
37
- repetition_penalty=1.1
38
- )
 
39
  response = tokenizer.decode(outputs[0], skip_special_tokens=True)
40
  response = response.split("[/INST]")[-1].strip()
41
  return {"question": q.question, "reponse": response}
 
6
 
7
  app = FastAPI(title="Assistant IA Education Marocaine")
8
 
 
9
  MODEL_NAME = "unsloth/mistral-7b-v0.3"
10
  LORA_NAME = "dohael/mistral-7b-education-maroc"
11
 
12
  tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
13
  model = AutoModelForCausalLM.from_pretrained(
14
  MODEL_NAME,
15
+ dtype=torch.float16,
16
+ device_map="cpu"
17
  )
18
+ model = PeftModel.from_pretrained(
19
+ model,
20
+ LORA_NAME,
21
+ is_trainable=False
22
+ )
23
+ model.eval()
24
 
25
  class Question(BaseModel):
26
  question: str
 
32
  @app.post("/ask")
33
  def ask(q: Question):
34
  prompt = f"<s>[INST] {q.question} [/INST]"
35
+ inputs = tokenizer(prompt, return_tensors="pt")
36
+ with torch.no_grad():
37
+ outputs = model.generate(
38
+ **inputs,
39
+ max_new_tokens=512,
40
+ temperature=0.7,
41
+ do_sample=True,
42
+ repetition_penalty=1.1
43
+ )
44
  response = tokenizer.decode(outputs[0], skip_special_tokens=True)
45
  response = response.split("[/INST]")[-1].strip()
46
  return {"question": q.question, "reponse": response}