jeffeux commited on
Commit
856d35d
1 Parent(s): d63c268
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -30,6 +30,7 @@ logging.info(C("[INFO] "f"device = {device}"))
30
  )
31
  def model_init():
32
 
 
33
 
34
 
35
  from transformers import GenerationConfig
@@ -65,15 +66,15 @@ try:
65
 
66
  # =================== INFERENCE ==================== #
67
  if prompt:
68
- st.balloons()
69
  with torch.no_grad():
70
  [texts_out] = model.generate(
71
  **tokenizer(
72
  prompt, return_tensors="pt",
73
 
74
  ).to(device),
75
- max_new_tokens=200,
76
- )
77
  output_text = tokenizer.decode(texts_out)
78
  st.balloons()
79
  st.markdown(output_text)
 
30
  )
31
  def model_init():
32
 
33
+ logging.info(C("[INFO] "f"Model init start!"))
34
 
35
 
36
  from transformers import GenerationConfig
 
66
 
67
  # =================== INFERENCE ==================== #
68
  if prompt:
69
+ st.title(prompt)
70
  with torch.no_grad():
71
  [texts_out] = model.generate(
72
  **tokenizer(
73
  prompt, return_tensors="pt",
74
 
75
  ).to(device),
76
+ max_new_tokens=100,
77
+ )
78
  output_text = tokenizer.decode(texts_out)
79
  st.balloons()
80
  st.markdown(output_text)