vilarin commited on
Commit
5312535
1 Parent(s): f2cc9dc

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -6
app.py CHANGED
@@ -36,7 +36,7 @@ tokenizer = AutoTokenizer.from_pretrained("THUDM/LongWriter-glm4-9b",trust_remot
36
 
37
 
38
  @spaces.GPU()
39
- def stream_chat(message: str, history: list, temperature: float, max_length: int):
40
  print(f'message is - {message}')
41
  print(f'history is - {history}')
42
  conversation = []
@@ -51,13 +51,12 @@ def stream_chat(message: str, history: list, temperature: float, max_length: int
51
  streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
52
 
53
  generate_kwargs = dict(
54
- max_length=max_length,
55
  streamer=streamer,
56
  do_sample=True,
57
  top_k=1,
58
  temperature=temperature,
59
- repetition_penalty=1.2,
60
- num_beams=1,
61
  )
62
  gen_kwargs = {**input_ids, **generate_kwargs}
63
 
@@ -88,11 +87,11 @@ with gr.Blocks(css=CSS) as demo:
88
  render=False,
89
  ),
90
  gr.Slider(
91
- minimum=128,
92
  maximum=32768,
93
  step=1,
94
  value=4096,
95
- label="Max Length",
96
  render=False,
97
  ),
98
  ],
 
36
 
37
 
38
  @spaces.GPU()
39
+ def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):
40
  print(f'message is - {message}')
41
  print(f'history is - {history}')
42
  conversation = []
 
51
  streamer = TextIteratorStreamer(tokenizer, timeout=60.0, skip_prompt=True, skip_special_tokens=True)
52
 
53
  generate_kwargs = dict(
54
+ max_new_tokens=max_new_tokens,
55
  streamer=streamer,
56
  do_sample=True,
57
  top_k=1,
58
  temperature=temperature,
59
+ repetition_penalty=1,
 
60
  )
61
  gen_kwargs = {**input_ids, **generate_kwargs}
62
 
 
87
  render=False,
88
  ),
89
  gr.Slider(
90
+ minimum=1024,
91
  maximum=32768,
92
  step=1,
93
  value=4096,
94
+ label="Max New Tokens",
95
  render=False,
96
  ),
97
  ],