Update app.py

app.py

@@ -1,64 +1,116 @@
+import os
+import requests
+import tempfile
+from pathlib import Path
+import secrets
+from PIL import Image
 import gradio as gr
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 
-# Use a pipeline as a high-level helper
-from transformers import pipeline
+# Set your Hugging Face API token
+HUGGING_FACE_API_KEY = os.getenv("HUGGING_FACE_API_KEY")
+math_messages = []
 
-messages = [
-    {"role": "user", "content": "Who are you?"},
-]
-pipe = pipeline("text-generation", model="Qwen/Qwen2.5-Math-1.5B")
-pipe(messages)
-# Load the model and tokenizer
-# model_name = "Qwen/Qwen2-Math-1.5B"
-# device = "cuda" if torch.cuda.is_available() else "cpu"
+# Function to process the image with Hugging Face API
+def process_image(image, shouldConvert=False):
+    global math_messages
+    math_messages = []  # Reset messages when a new image is uploaded
 
-# model = AutoModelForCausalLM.from_pretrained(
-#     model_name,
-#     torch_dtype="auto",
-#     device_map="auto"
-# ).to(device)
+    uploaded_file_dir = os.environ.get("GRADIO_TEMP_DIR") or str(Path(tempfile.gettempdir()) / "gradio")
+    os.makedirs(uploaded_file_dir, exist_ok=True)
+    
+    name = f"tmp{secrets.token_hex(20)}.jpg"
+    filename = os.path.join(uploaded_file_dir, name)
+    
+    # Save the uploaded image
+    if shouldConvert:
+        new_img = Image.new('RGB', (image.width, image.height), (255, 255, 255))
+        new_img.paste(image, (0, 0), mask=image)
+        image = new_img
+    image.save(filename)
+    
+    # Use Hugging Face API for image captioning
+    with open(filename, "rb") as img_file:
+        response = requests.post(
+            "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-base",
+            headers={"Authorization": f"Bearer {HUGGING_FACE_API_KEY}"},
+            files={"file": img_file}
+        )
+    os.remove(filename)  # Clean up temp file
+    
+    caption = response.json().get("generated_text", "No description available.")
+    return caption
 
-# tokenizer = AutoTokenizer.from_pretrained(model_name)
+# Function for getting math responses from Hugging Face's text generation API
+def get_math_response(image_description, user_question):
+    global math_messages
+    if not math_messages:
+        math_messages.append({"role": "system", "content": "You are a helpful math assistant."})
+    
+    # Prepare the query content
+    content = f"Image description: {image_description}\n\n" if image_description else ""
+    query = f"{content}User question: {user_question}"
+    math_messages.append({"role": "user", "content": query})
+    
+    # Make the text generation call
+    payload = {
+        "inputs": query,
+        "parameters": {"max_length": 100, "temperature": 0.7},
+    }
+    response = requests.post(
+        "https://api-inference.huggingface.co/models/gpt2",
+        headers={"Authorization": f"Bearer {HUGGING_FACE_API_KEY}"},
+        json=payload
+    )
+    
+    answer = response.json().get("generated_text", "Sorry, I couldn't generate a response.")
+    yield answer
+    math_messages.append({"role": "assistant", "content": answer})
 
-# # Define a function for Gradio to handle user input
-# def solve_math(prompt):
-#     messages = [
-#         {"role": "system", "content": "You are a helpful assistant."},
-#         {"role": "user", "content": prompt}
-#     ]
-#     text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-#     model_inputs = tokenizer([text], return_tensors="pt").to(device)
+def math_chat_bot(image, sketchpad, question, state):
+    current_tab_index = state["tab_index"]
+    image_description = None
+    
+    # Check for uploaded image
+    if current_tab_index == 0 and image:
+        image_description = process_image(image)
+    elif current_tab_index == 1 and sketchpad and sketchpad["composite"]:
+        image_description = process_image(sketchpad["composite"], True)
+    
+    # Get response from the text generation API
+    yield from get_math_response(image_description, question)
 
-#     generation_config = GenerationConfig(
-#         do_sample=False,  # For greedy decoding
-#         max_new_tokens=512
-#     )
+css = """
+#qwen-md .katex-display { display: inline; }
+#qwen-md .katex-display>.katex { display: inline; }
+#qwen-md .katex-display>.katex>.katex-html { display: inline; }
+"""
 
-#     generated_ids = model.generate(
-#         **model_inputs,
-#         generation_config=generation_config
-#     )
+def tabs_select(e: gr.SelectData, _state):
+    _state["tab_index"] = e.index
 
-#     # Remove the input tokens from the output
-#     generated_ids = [
-#         output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
-#     ]
-
-#     # Decode the generated output and return the result
-#     response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
-#     return response
-
-# # Create the Gradio interface
-# iface = gr.Interface(
-#     fn=solve_math,  # Function to call
-#     inputs="text",  # Text input for the user prompt
-#     outputs="text", # Text output for the model's response
-#     title="Math Solver",  # App title
-#     description="Provide a math problem and the model will solve it."
-# )
-
-# Launch the app
-if __name__ == "__main__":
-    iface.launch()
+# Create Gradio interface
+with gr.Blocks(css=css) as demo:
+    gr.HTML("""<center><font size=8>📖 Math Assistant Demo</center>""")
+    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", elem_id="qwen-md")
+        submit_btn.click(
+            fn=math_chat_bot,
+            inputs=[*input_image, input_sketchpad, input_text, state],
+            outputs=output_md
+        )
+demo.launch()