To Chat Interface

Quick PR still not finished to make it a chat interface instead! Almost done, just history logic to be done, will do later 👍

app.py

@@ -10,7 +10,7 @@ from mistral_inference.transformer import Transformer
 from mistral_inference.generate import generate
 
 from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-from mistral_common.protocol.instruct.messages import UserMessage, TextChunk, ImageURLChunk
+from mistral_common.protocol.instruct.messages import UserMessage, AssistantMessage, TextChunk, ImageURLChunk
 from mistral_common.protocol.instruct.request import ChatCompletionRequest
 
 models_path = Path.home().joinpath('pixtral', 'Pixtral')
@@ -29,9 +29,20 @@ def image_to_base64(image_path):
     return f"data:image/jpeg;base64,{encoded_string}"
 
 @spaces.GPU(duration=30)
-def run_inference(image_url, prompt):
-    base64 = image_to_base64(image_url)
-    completion_request = ChatCompletionRequest(messages=[UserMessage(content=[ImageURLChunk(image_url=base64), TextChunk(text=prompt)])])
+def run_inference(message, history):
+    print(message)
+    print(history)
+
+    ## to be fixed
+    messages = []
+    for couple in history:
+        messages.append(UserMessage(content = [ImageURLChunk(image_url=image_to_base64(file["path"])) for file in couple[0][0]]+[TextChunk(text=couple[0][1])]))
+        messages.append(AssistantMessage(content = couple[1]))
+    ##
+        
+    messages.append(UserMessage(content = [ImageURLChunk(image_url=image_to_base64(file["path"])) for file in message["files"]]+[TextChunk(text=message["text"])]))
+    
+    completion_request = ChatCompletionRequest(messages=messages)
     
     encoded = tokenizer.encode_chat_completion(completion_request)
     
@@ -40,26 +51,7 @@ def run_inference(image_url, prompt):
     
     out_tokens, _ = generate([tokens], model, images=[images], max_tokens=512, temperature=0.45, eos_id=tokenizer.instruct_tokenizer.tokenizer.eos_id)
     result = tokenizer.decode(out_tokens[0])
-    return [[prompt, result]]
-
-with gr.Blocks() as demo:
-    with gr.Row():
-        image_box = gr.Image(type="filepath")
-   
-        chatbot = gr.Chatbot(
-            scale = 2,
-            height=750
-        )
-    text_box = gr.Textbox(
-            placeholder="Enter your text and press enter, or upload an image.",
-            container=False,
-        )
-
-
-    btn = gr.Button("Submit")
-    clicked = btn.click(run_inference,
-                        [image_box,text_box],
-                        chatbot
-                        )
+    return result
 
+demo = gr.ChatInterface(fn=run_inference, title="Pixtral 12B", multimodal=True)
 demo.queue().launch()