Update app.py

app.py

@@ -35,6 +35,7 @@ model = AutoModelForCausalLM.from_pretrained(
 
 tokenizer = AutoTokenizer.from_pretrained(MODELS,trust_remote_code=True)
 
+
 @spaces.GPU
 def stream_chat(message: str, history: list, temperature: float, max_new_tokens: int):
  print(f'message is - {message}')
@@ -48,17 +49,18 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
  
  input_ids = tokenizer.apply_chat_template(conversation, tokenize=True, add_generation_prompt=True, return_tensors="pt", return_dict=True).to(model.device)
  
- streamer = TextIteratorStreamer(tokenizer, timeout=60, skip_prompt=True, skip_special_tokens=True)
+ # streamer = TextIteratorStreamer(tokenizer, timeout=60, skip_prompt=True, skip_special_tokens=True)
 
  generate_kwargs = dict(
  input_ids=input_ids, 
- streamer=streamer,
- max_length=max_new_tokens, 
- do_sample=True, 
+ max_length=2500, 
+ max_new_tokens=max_new_tokens,
+ do_sample=True,
+ top_k=1,
  temperature=temperature,
  repetition_penalty=1.2,
  )
- 
+'''
  thread = Thread(target=model.generate, kwargs=generate_kwargs)
  thread.start()
 
@@ -66,6 +68,13 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
  for new_text in streamer:
  buffer[-1][1] += new_text
  yield buffer
+''' 
+ with torch.no_grad():
+ outputs = model.generate(**inputs, **gen_kwargs)
+ outputs = outputs[:, inputs['input_ids'].shape[1]:]
+ results = tokenizer.decode(outputs[0], skip_special_tokens=True)
+ return results
+