Initial commit

app.py

@@ -0,0 +1,170 @@
+"""
+File: app.py
+
+Description: Visual question answering (with pixtral-12b-2409)
+
+Author: Didier Guillevic
+Date: 2024-09-29
+"""
+import logging
+logger = logging.getLogger(__name__)
+
+import gradio as gr
+import PIL
+import base64
+from io import BytesIO
+
+import os
+from mistralai import Mistral
+
+#
+# Mistral AI client
+#
+api_key = os.environ["MISTRAL_API_KEY"]
+client = Mistral(api_key=api_key)
+model_id = "pixtral-12b-2409"
+
+#
+# Utils
+#
+def image_to_base64(filepath):
+    file_format = filepath.split('.')[-1].lower()
+    with open(filepath, "rb") as image_file:
+        base64_string = base64.b64encode(image_file.read()).decode('utf-8')
+    return base64_string, file_format
+
+def pil_to_base64(pil_image, image_format='jpeg'):  # Default to 'png' if format not known
+    # Save the PIL image to a BytesIO object in the specified format
+    buffered = BytesIO()
+    pil_image.save(buffered, format=image_format)
+    
+    # Convert the image to a base64 string
+    base64_string = base64.b64encode(buffered.getvalue()).decode('utf-8')
+    
+    return base64_string, image_format
+
+
+#
+# Generate a response
+#
+def answer_question(text, image, temperature=0.0, max_tokens=1_024):
+    # Convert image to base64 string
+    #base64_string, file_format = image_to_base64(image)
+    base64_string, file_format = pil_to_base64(image)
+
+    messages = [
+        {
+            "role": "user",
+            "content": [
+                {
+                    "type": "text",
+                    "text": text
+                },
+                {
+                    "type": "image_url",
+                    "image_url": f"data:image/{file_format};base64,{base64_string}" 
+                }
+            ]
+        }
+    ]
+    
+    chat_response = client.chat.complete(
+        model=model_id,
+        messages=messages,
+        temperature=temperature,
+        max_tokens=max_tokens
+    )
+
+    return chat_response.choices[0].message.content
+
+
+def clear_all():
+    # Clear question, image, output_text
+    return ('', None, '')
+
+
+#
+# User interface
+#
+with gr.Blocks() as demo:
+    
+    gr.Markdown("""
+        ## Visual question answering / image captioning
+    """)
+    with gr.Row():
+        with gr.Column():
+            question = gr.Textbox(
+                placeholder="Ask anything about the image..",
+                lines=2,
+                render=True)
+            image = gr.Image(type="pil") # return type: 'numpy', 'pil', 'filepath'
+            with gr.Row():
+                temperature = gr.Slider(
+                    label="Temperature",
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=0.0,
+                    step=0.1
+                )
+                max_tokens = gr.Slider(
+                    label="Max tokens",
+                    minimum=128,
+                    maximum=2_048,
+                    value=1_024,
+                    step=128
+                )
+        with gr.Column():
+            output_text = gr.Textbox(
+                lines=10,
+                label="Pixtral 12B",
+                render=True
+            )
+    with gr.Row():
+        clear_btn = gr.Button("Clear", variant="secondary")
+        submit_btn = gr.Button(value="Submit", variant="primary")
+
+    # Examples
+    examples = gr.Examples(
+        [
+            [
+                'Can you describe this image?', './sample_ID.jpeg'
+            ],
+            [
+                'In which city and country was  this picture taken?',
+                './16-park-ave-milton-park-1024x683.jpg'
+            ],
+            [   (
+                    'Please extract the identification information present in the image. '
+                    'Format the response in a JSON format such as "{name_of_element : [value]}".'
+                ),
+                './driver_license.png'
+            ],
+        ],
+        inputs=[question, image],
+        outputs=[output_text,],
+        fn=answer_question,
+        cache_examples=False,
+        label="Examples"
+    )
+
+    # Click actions
+    clear_btn.click(
+        clear_all,
+        inputs=[],
+        outputs=[question, image, output_text]
+    )
+    submit_btn.click(
+        fn=answer_question,
+        inputs=[question, image, temperature, max_tokens],
+        outputs=[output_text,]
+    )
+
+    # Documentation
+    with gr.Accordion("Documentation", open=False):
+        gr.Markdown("""
+            - model: Serving pixtral-12b (https://mistral.ai/news/pixtral-12b/)
+            - temperature: 0.0=little variable in output, 1.0=lots of variability in output
+            - max_tokens: maximum number of tokens to generate as output (1 word approx 1.5 tokens)
+        """)
+
+demo.launch()