Update app.py

app.py

@@ -32,7 +32,7 @@ model = AutoModelForCausalLM.from_pretrained(
         trust_remote_code=True,
         ).eval()
 
-tokenizer = AutoTokenizer.from_pretrained("THUDM/LongWriter-glm4-9b",trust_remote_code=True)
+tokenizer = AutoTokenizer.from_pretrained("THUDM/LongWriter-glm4-9b",trust_remote_code=True, use_fast=False)
 
 class StopOnTokens(StoppingCriteria):
     def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
@@ -56,7 +56,7 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
     print(f"Conversation is -\n{conversation}")
     stop = StopOnTokens()
 
-    input_ids = tokenizer.build_chat_input(message, history=conversation, role='user').input_ids.to(model.device)
+    input_ids = tokenizer.build_chat_input(message, history=conversation, role='user').input_ids.to(next(model.parameters()).device)
     #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.0, skip_prompt=True, skip_special_tokens=True)
     eos_token_id = [tokenizer.eos_token_id, tokenizer.get_command("<|user|>"),
@@ -64,8 +64,8 @@ def stream_chat(message: str, history: list, temperature: float, max_new_tokens:
 
     generate_kwargs = dict(
         input_ids=input_ids,
-        max_new_tokens=max_new_tokens,
         streamer=streamer,
+        max_new_tokens=max_new_tokens,
         do_sample=True,
         top_k=1,
         temperature=temperature,