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
+