import gradio as gr from transformers import AutoModelForCausalLM, AutoProcessor, GenerationConfig from PIL import Image import torch import spaces import json import re import deepl # Load the processor and model processor = AutoProcessor.from_pretrained( 'allenai/Molmo-7B-D-0924', trust_remote_code=True, torch_dtype='auto', device_map='auto' ) model = AutoModelForCausalLM.from_pretrained( 'allenai/Molmo-7B-D-0924', trust_remote_code=True, torch_dtype='auto', device_map='auto' ) @spaces.GPU() def wrap_json_in_markdown(text): result = [] stack = [] json_start = None in_json = False i = 0 while i < len(text): char = text[i] if char in ['{', '[']: if not in_json: json_start = i in_json = True stack.append(char) else: stack.append(char) elif char in ['}', ']'] and in_json: if not stack: # Unbalanced bracket, reset in_json = False json_start = None else: last = stack.pop() if (last == '{' and char != '}') or (last == '[' and char != ']'): # Mismatched brackets in_json = False json_start = None if in_json and not stack: # Potential end of JSON json_str = text[json_start:i+1] try: # Try to parse the JSON to ensure it's valid parsed = json.loads(json_str) # Wrap in Markdown code block wrapped = f"\n```json\n{json.dumps(parsed, indent=4)}\n```\n" result.append(text[:json_start]) # Append text before JSON result.append(wrapped) # Append wrapped JSON text = text[i+1:] # Update the remaining text i = -1 # Reset index except json.JSONDecodeError: # Not valid JSON, continue searching pass in_json = False json_start = None i += 1 result.append(text) # Append any remaining text return ''.join(result) def decode_unicode_sequences(unicode_seq): """ Decodes a sequence of Unicode escape sequences (e.g., \\u4F60\\u597D) to actual characters. Args: unicode_seq (str): A string containing Unicode escape sequences. Returns: str: The decoded Unicode string. """ # Regular expression to find \uXXXX unicode_escape_pattern = re.compile(r'\\u([0-9a-fA-F]{4})') # Function to replace each \uXXXX with the corresponding character def replace_match(match): hex_value = match.group(1) return chr(int(hex_value, 16)) # Decode all \uXXXX sequences decoded = unicode_escape_pattern.sub(replace_match, unicode_seq) return decoded def is_mandarin(text): """ Detects if the given text is in Mandarin using Unicode ranges. Args: text (str): The text to check. Returns: bool: True if the text contains Chinese characters, False otherwise. """ # Chinese Unicode ranges for char in text: if '\u4e00' <= char <= '\u9fff': return True return False def translate_to_english_deepl(text, api_key): """ Translates Mandarin text to English using DeepL API. Args: text (str): The Mandarin text to translate. api_key (str): Your DeepL API authentication key. Returns: str: The translated English text. """ url = "https://api.deepl.com/v2/translate" params = { "auth_key": api_key, "text": text, "source_lang": "ZH", "target_lang": "EN" } # try: # response = requests.post(url, data=params) # response.raise_for_status() # result = response.json() # return result['translations'][0]['text'] # except requests.exceptions.RequestException as e: # print(f"DeepL Translation error: {e}") # return text # Return the original text if translation fails # auth_key = api_key # Replace with your key # translator = deepl.Translator(auth_key) # result = translator.translate_text("Hello, world!", target_lang="FR") # print(result.text) # "Bonjour, le monde !" try: auth_key = api_key # Replace with your key translator = deepl.Translator(auth_key) result = translator.translate_text(text, source_lang="ZH", target_lang="EN-US") # print(result.text) return result.text except requests.exceptions.RequestException as e: print(f"DeepL Translation error: {e}") return text # Return the original text if translation fails def process_text_deepl(input_string, api_key): """ Processes the input string to find Unicode escape sequences representing Mandarin words, translates them to English using DeepL, and replaces them accordingly. Args: input_string (str): The original string containing Unicode escape sequences. api_key (str): Your DeepL API authentication key. Returns: str: The processed string with translations where applicable. """ # Regular expression to find groups of consecutive \uXXXX sequences unicode_word_pattern = re.compile(r'(?:\\u[0-9a-fA-F]{4})+') # Function to process each matched Unicode word def process_match(match): unicode_seq = match.group(0) decoded_word = decode_unicode_sequences(unicode_seq) if is_mandarin(decoded_word): translated = translate_to_english_deepl(decoded_word, api_key) return f"{translated} ({decoded_word})" else: # If not Mandarin, return the original sequence return unicode_seq # Substitute all matched Unicode words with their translations if applicable processed_string = unicode_word_pattern.sub(process_match, input_string) return processed_string def process_image_and_text(image, text): # Process the image and text inputs = processor.process( images=[Image.fromarray(image)], text=text ) # Move inputs to the correct device and make a batch of size 1 inputs = {k: v.to(model.device).unsqueeze(0) for k, v in inputs.items()} # Generate output output = model.generate_from_batch( inputs, GenerationConfig(max_new_tokens=1024, stop_strings="<|endoftext|>"), tokenizer=processor.tokenizer ) # Only get generated tokens; decode them to text generated_tokens = output[0, inputs['input_ids'].size(1):] generated_text = processor.tokenizer.decode(generated_tokens, skip_special_tokens=True) generated_text_w_json_wrapper = wrap_json_in_markdown(generated_text) generated_text_w_unicode_mdn = process_text_deepl(generated_text_w_json_wrapper, "a5b1749b-7112-4c2d-81a3-33ea18478bb4:fx") return generated_text_w_json_wrapper def chatbot(image, text, history): if image is None: return history + [("Please upload an image first.", None)] response = process_image_and_text(image, text) history.append({"role": "user", "content": text}) history.append({"role": "assistant", "content": response}) return history # Define the Gradio interface with gr.Blocks() as demo: gr.Markdown("# Image Chatbot with Molmo-7B-D-0924") with gr.Row(): image_input = gr.Image(type="numpy") chatbot_output = gr.Chatbot(type="messages") text_input = gr.Textbox(placeholder="Ask a question about the image...") submit_button = gr.Button("Submit") state = gr.State([]) submit_button.click( chatbot, inputs=[image_input, text_input, state], outputs=[chatbot_output] ) text_input.submit( chatbot, inputs=[image_input, text_input, state], outputs=[chatbot_output] ) demo.launch()