Update app.py

app.py

@@ -48,7 +48,11 @@ Assistant: girlfriend;mother;father;friend
 torch.set_grad_enabled(False)
 model_name = "TheBloke/OpenHermes-2.5-Mistral-7B-GPTQ"
 pipe = pipeline("text-generation", model=model_name, device='cuda')
-generate_kwargs = {'max_new_tokens': 20}
+generate_kwargs = {
+    'max_new_tokens': 20, 
+    'temperature': 0.8, 
+    'repetition_penalty': 1.1
+}
 
 def past_kv_to_device(past_kv, device, dtype):
     return tuple((torch.tensor(k).to(device).to(dtype), torch.tensor(v).to(device).to(dtype)) for k, v in past_kv)
@@ -84,14 +88,15 @@ def generate(text, past_key_values):
         prompt_format.format(system_message=system_prompt, prompt=text), **cur_generate_kwargs
     )[0]['generated_text']
     print(response)
-    return response[-1]['content']
+    return response.split('<|im_start|>assistant')[1]
+    # return response[-1]['content']
 
     
 if __name__ == "__main__":
     with torch.no_grad():
-        past_key_values = set_past_key_values()
+        # past_key_values = set_past_key_values()
         demo = gr.Interface(
-            partial(generate, past_key_values=past_key_values), 
+            partial(generate, past_key_values=None), 
             inputs="textbox", outputs="textbox"
         )
         demo.launch()