Upload app.py

app.py

@@ -0,0 +1,92 @@
+import torch
+from transformers import LlavaForConditionalGeneration, BitsAndBytesConfig, AutoProcessor
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+import requests
+from PIL import Image
+import requests
+import gradio as gr
+
+
+# Load translation model and tokenizer
+translate_model_name = "facebook/mbart-large-50-many-to-many-mmt"
+translate_model = MBartForConditionalGeneration.from_pretrained(translate_model_name)
+tokenizer = MBart50TokenizerFast.from_pretrained(translate_model_name)
+
+# load the base model in 4 bit quantized
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+)
+
+# finetuned model adapter path (Hugging Face Hub)
+model_id = 'somnathsingh31/llava-1.5-7b-hf-ft-merged_model'
+
+# merge the models
+merged_model = LlavaForConditionalGeneration.from_pretrained(model_id,
+                                                           quantization_config=quantization_config,
+                                                           torch_dtype=torch.float16)
+
+# create processor from base model
+processor = AutoProcessor.from_pretrained("llava-hf/llava-1.5-7b-hf")
+
+# function to translate
+def translate(text, source_lang, target_lang):
+    # Set source language
+    tokenizer.src_lang = source_lang
+
+    # Encode the text
+    encoded_text = tokenizer(text, return_tensors="pt")
+
+    # Force target language token
+    forced_bos_token_id = tokenizer.lang_code_to_id[target_lang]
+
+    # Generate the translation
+    generated_tokens = translate_model.generate(**encoded_text, forced_bos_token_id=forced_bos_token_id)
+
+    # Decode the translation
+    translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
+
+    return translation
+
+
+# function for making inference
+def ask_vlm(hindi_input_text, image):
+    # translate from Hindi to English
+    prompt_eng = translate(hindi_input_text, "hi_IN", "en_XX")
+    prompt = "USER: <image>\n" + prompt_eng + " ASSISTANT:"
+
+    # If image is uploaded, open the image from bytes, else open from URL
+    if hasattr(image, 'read'):
+        image = Image.open(image)
+    else:
+        image = Image.open(requests.get(image, stream=True).raw)
+    
+    inputs = processor(text=prompt, images=image, return_tensors="pt")
+    generate_ids = merged_model.generate(**inputs, max_new_tokens=250)
+    decoded_response = processor.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
+    assistant_index = decoded_response.find("ASSISTANT:")
+
+    # Extract text after "ASSISTANT:"
+    if assistant_index != -1:
+        text_after_assistant = decoded_response[assistant_index + len("ASSISTANT:"):]
+        # Remove leading and trailing whitespace
+        text_after_assistant = text_after_assistant.strip()
+    else:
+        text_after_assistant = None
+
+    hindi_output_text = translate(text_after_assistant, "en_XX", "hi_IN")
+    return hindi_output_text
+
+# Define Gradio interface
+input_image = gr.inputs.Image(type="pil", label="Input Image (Upload or URL)")
+input_question = gr.inputs.Textbox(lines=2, label="Question (Hindi)")
+output_text = gr.outputs.Textbox(label="Response (Hindi)")
+
+# Create Gradio app
+gr.Interface(fn=ask_vlm, inputs=[input_question, input_image], outputs=output_text, title="Image and Text-based Dialogue System", description="Enter a question in Hindi and an image, either by uploading or providing URL, and get a response in Hindi.").launch()
+
+
+if __name__ ==  '__main__':
+    image_url = 'https://images.metmuseum.org/CRDImages/ad/original/138425.jpg'
+    user_query = 'यह किस प्रकार की कला है? विस्तार से बताइये'
+    output = ask_vlm(user_query, image_url)
+    print('Output:\n', output)