Update app.py

app.py

@@ -73,16 +73,27 @@ def chat_llm_jp_v2(message: str,
     
     streamer = TextIteratorStreamer(tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True)
 
-    output = model.generate(
-        input_ids,
+    generate_kwargs = dict(
+        input_ids= input_ids,
+        streamer=streamer,
         max_new_tokens=max_new_tokens,
         do_sample=True,
-        top_p=0.95,
         temperature=temperature,
-        repetition_penalty=1.05,
-    )[0]
-                 
-    return tokenizer.decode(output)
+        top_p=0.95,
+        repetition_penalty=1.1,
+    )
+    # This will enforce greedy generation (do_sample=False) when the temperature is passed 0, avoiding the crash.             
+    if temperature == 0:
+        generate_kwargs['do_sample'] = False
+        
+    t = Thread(target=model.generate, kwargs=generate_kwargs)
+    t.start()
+
+    outputs = []
+    for text in streamer:
+        outputs.append(text)
+        print(outputs)
+        yield "".join(outputs)
         
 
 # Gradio block