Update app.py

app.py

@@ -57,20 +57,13 @@ def generate_text(messages):
         model = AutoModelForCausalLM.from_pretrained(
                 model_id, token=huggingface_token ,torch_dtype=dtype,device_map=device
             )
-        text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device ) #pipeline has not to(device)
+        text_generator = pipeline("text-generation", model=model, tokenizer=tokenizer,torch_dtype=dtype,device_map=device ,stream=True) #pipeline has not to(device)
     result = text_generator(messages, max_new_tokens=256, do_sample=True, temperature=0.7)
 
-    generated_output = result[0]["generated_text"]
-    if isinstance(generated_output, list):
-        for message in reversed(generated_output):
-            if message.get("role") == "assistant":
-                content= message.get("content", "No content found.")
-                return content
-            
-        return "No assistant response found."
-    else:
-        return "Unexpected output format."
-
+    generated_output = ""
+    for token in result:
+        generated_output += token["generated_token"]
+        yield generated_output
 
 
 def call_generate_text(message, history):
@@ -80,12 +73,12 @@ def call_generate_text(message, history):
    
     messages = history+[{"role":"user","content":message}]
     try:
-        text = generate_text(messages)
-        return text
+        
+        for text in generate_text(messages):
+            yield text
     except RuntimeError  as e:
         print(f"An unexpected error occurred: {e}")
-       
-    return ""
+        yield ""
 
 demo = gr.ChatInterface(call_generate_text,type="messages")