analist commited on
Commit
c745c2b
·
verified ·
1 Parent(s): 591d0d9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -2
app.py CHANGED
@@ -12,8 +12,8 @@ def load_model():
12
  tokenizer = AutoTokenizer.from_pretrained(model_name)
13
  model = AutoModelForCausalLM.from_pretrained(
14
  model_name,
15
- torch_dtype=torch.float16,
16
- device_map="auto"
17
  )
18
  return model, tokenizer
19
 
@@ -79,6 +79,7 @@ def main():
79
  with st.spinner("Réflexion en cours..."):
80
  prompt = format_prompt(question)
81
  response = generate_response(prompt, model, tokenizer)
 
82
  st.markdown(response)
83
  st.session_state.messages.append({"role": "assistant", "content": response})
84
 
 
12
  tokenizer = AutoTokenizer.from_pretrained(model_name)
13
  model = AutoModelForCausalLM.from_pretrained(
14
  model_name,
15
+ torch_dtype=torch.float32,
16
+ device_map="cpu"
17
  )
18
  return model, tokenizer
19
 
 
79
  with st.spinner("Réflexion en cours..."):
80
  prompt = format_prompt(question)
81
  response = generate_response(prompt, model, tokenizer)
82
+ response = response.replace('<think>', '')
83
  st.markdown(response)
84
  st.session_state.messages.append({"role": "assistant", "content": response})
85