Update app.py

app.py

@@ -233,13 +233,18 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             choice, contents = mode_load(file_name)
 
             if choice == "image":
-                input_text = message['text']
+                #input_text = message['text']
+                conversation.append({"role": "user", "image": contents, "content": message['text']})
             elif choice == "doc":
-                input_text = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
+                #input_text = contents + "\n\n\n" + "{} files uploaded.\n".format(1) + message['text']
+                format_msg = contents + "\n\n\n" + "{} files uploaded.\n" + message['text']
+                conversation.append({"role": "user", "content": format_msg})
         else:
-            input_text = message['text']
+            #input_text = message['text']
+            contents = None
+            conversation.append({"role": "user", "content": message['text']})
 
-        conversation = [{"role": "user", "content": input_text}]
+        # conversation = [{"role": "user", "content": input_text}]
         input_ids = tokenizer(conversation[-1]['content'], return_tensors="pt").to(model.device)
 
         generate_kwargs = dict(
@@ -257,8 +262,11 @@ def simple_chat(message: dict, temperature: float = 0.8, max_length: int = 4096,
             generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
 
         # Process to remove any prefix or unwanted prompt
-        text_original = input_text.strip()
-        results_text = generated_text[len(text_original):].strip()
+        # text_original = input_text.strip()
+        # results_text = generated_text[len(text_original):].strip()
+
+        text_original = message['text'].strip()
+        results_text = generated_text.replace(text_original, "").strip()
 
         print("  ")
         print("------")