Fix add logs

app.py

@@ -21,8 +21,8 @@ async def handle(system_input: str, user_input: str):
             buffers[i] += str(outputs[i])
 
         yield list(buffers) + ["", ""]
-    yield list(buffers) + [generate_openllama(system_input, user_input),
-                           generate_bloom(system_input, user_input)]
+    yield list(buffers) + [(openllama_generation := generate_openllama(system_input, user_input)), '']
+    yield list(buffers) + [openllama_generation, generate_bloom(system_input, user_input)]
 
 
 with gr.Blocks() as demo:

generators.py

@@ -42,7 +42,7 @@ async def generate_gpt2(system_input, user_input):
     output = await query_llm({
         "inputs": (inputs:=f"{system_input}\n{user_input}"),
     }, "openai-community/gpt2")
-    yield output[0]["generated_text"].replace(inputs, '')
+    yield output[0]["generated_text"]
 
 
 async def generate_llama2(system_input, user_input):
@@ -67,7 +67,7 @@ def generate_openllama(system_input, user_input):
     model = LlamaForCausalLM.from_pretrained(
         model_path, torch_dtype=torch.float16, device_map='cuda',
     )
-    # model = model.to("cuda")
+    print('model openllama loaded')
     input_text = f"{system_input}\n{user_input}"
     input_ids = tokenizer(input_text, return_tensors="pt").input_ids.to("cuda")
     output = model.generate(input_ids, max_length=128)