Update app.py

app.py

@@ -1,17 +1,162 @@
-import gradio as gr
-from transformers import pipeline
+from langchain.prompts import PromptTemplate
+from langchain_huggingface import HuggingFaceEndpoint
+from PIL import Image
+import os
+import secrets
+from pathlib import Path
+import tempfile
 
-def generate_response(text):
-  pipe = pipeline("text-generation", model="text-davinci-003")
-  response = pipe(text, max_length=512)[0]['generated_text']
-  return response
+# Initialize the Hugging Face BLIP model
+image_captioning_model = HuggingFaceEndpoint(
+    endpoint_url="https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-base",
+    huggingfacehub_api_token=os.getenv("HUGGING_FACE_API"),  # Ensure you set this in your environment
+    temperature=0.7,
+    max_new_tokens=1024,
+)
+math_llm=HuggingFaceEndpoint(
+    endpoint_url="https://api-inference.huggingface.co/models/meta-llama/Llama-3.2-3B-Instruct",
+    huggingfacehub_api_token=os.getenv("HUGGING_FACE_API"),  # Ensure you set this in your environment
+    temperature=0.7,
+    max_new_tokens=1024,)
+# Function to process the image
+def process_image(image, shouldConvert=False):
+    # Ensure temporary directory exists
+    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(
+        Path(tempfile.gettempdir()) / "gradio"
+    )
+    os.makedirs(uploaded_file_dir, exist_ok=True)
+    
+    # Save the uploaded image
+    name = f"tmp{secrets.token_hex(20)}.jpg"
+    filename = os.path.join(uploaded_file_dir, name)
+    if shouldConvert:
+        # Convert image to RGB mode if it contains transparency
+        new_img = Image.new("RGB", size=(image.width, image.height), color=(255, 255, 255))
+        new_img.paste(image, (0, 0), mask=image)
+        image = new_img
+    image.save(filename)
 
-with gr.Blocks() as demo:
-  gr.Markdown("## Math Reasoning Chatbot")
-  txt = gr.Textbox(label="Enter your math problem")
-  btn = gr.Button("Submit")
-  output = gr.Textbox()
+    # Define a PromptTemplate for text instruction
+    template = """
+    You are a helpful AI assistant.
+    Please describe the math-related content in this image, ensuring that any LaTeX formulas are correctly transcribed. 
+    Non-mathematical details do not need to be described.
 
-  btn.click(fn=generate_response, inputs=txt, outputs=output)
+    Image Path: {image}
+    """
+    prompt_template = PromptTemplate(
+        input_variables=["image"],  # Dynamically insert the image path
+        template=template
+    )
 
+    # Create the text instruction by rendering the prompt template
+    prompt = prompt_template.format(image=f"file://{filename}")
+
+    # Use the model with both the image and the generated prompt
+    with open(filename, "rb") as img_file:
+        response = image_captioning_model({
+            "inputs": {
+                "image": img_file,
+                "text": prompt
+            }
+        })
+
+    # Return the model's response
+    return response
+
+def get_math_response(image_description, user_question):
+    template = """
+    You are a helpful AI assistant specialized in solving math reasoning problems. 
+    Analyze the following question carefully and provide a step-by-step explanation along with the answer.
+    Image description : {image_description}
+    Question: {user_question}?
+    """
+
+    prompt_template = PromptTemplate(
+        input_variables=["user_question","image_description"],  # Define the placeholder(s) in the template
+        template=template
+    )
+    formatted_prompt = prompt_template.format(user_question=user_question, image_description=image_description)
+
+    # Pass the formatted prompt to the model
+    response = math_llm(formatted_prompt)
+
+    # Print the response
+    yield response
+
+def math_chat_bot(image, sketchpad, question, state):
+    current_tab_index = state["tab_index"]
+    image_description = None
+    # Upload
+    if current_tab_index == 0:
+        if image is not None:
+            image_description = process_image(image)
+    # Sketch
+    elif current_tab_index == 1:
+        print(sketchpad)
+        if sketchpad and sketchpad["composite"]:
+            image_description = process_image(sketchpad["composite"], True)
+    yield from get_math_response(image_description, question)
+css = """
+#qwen-md .katex-display { display: inline; }
+#qwen-md .katex-display>.katex { display: inline; }
+#qwen-md .katex-display>.katex>.katex-html { display: inline; }
+"""
+def tabs_select(e: gr.SelectData, _state):
+    _state["tab_index"] = e.index
+
+with gr.Blocks(css=css) as demo:
+
+    state = gr.State({"tab_index": 0})
+    with gr.Row():
+        with gr.Column():
+            with gr.Tabs() as input_tabs:
+                with gr.Tab("Upload"):
+                    input_image = gr.Image(type="pil", label="Upload"),
+                with gr.Tab("Sketch"):
+                    input_sketchpad = gr.Sketchpad(type="pil", label="Sketch", layers=False)
+            input_tabs.select(fn=tabs_select, inputs=[state])
+            input_text = gr.Textbox(label="input your question")
+            with gr.Row():
+                with gr.Column():
+                    clear_btn = gr.ClearButton(
+                        [*input_image, input_sketchpad, input_text])
+                with gr.Column():
+                    submit_btn = gr.Button("Submit", variant="primary")
+        with gr.Column():
+            output_md = gr.Markdown(label="answer",
+                                    latex_delimiters=[{
+                                        "left": "\\(",
+                                        "right": "\\)",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{equation\}",
+                                        "right": "\\end\{equation\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{align\}",
+                                        "right": "\\end\{align\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{alignat\}",
+                                        "right": "\\end\{alignat\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{gather\}",
+                                        "right": "\\end\{gather\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\begin\{CD\}",
+                                        "right": "\\end\{CD\}",
+                                        "display": True
+                                    }, {
+                                        "left": "\\[",
+                                        "right": "\\]",
+                                        "display": True
+                                    }],
+                                    elem_id="qwen-md")
+        submit_btn.click(
+            fn=math_chat_bot,
+            inputs=[*input_image, input_sketchpad, input_text, state],
+            outputs=output_md)
 demo.launch()