Added the transformer to translate text from English to Spanish

app.py

@@ -5,7 +5,13 @@ import gradio as gr
 import pandas as pd
 from datasets import load_dataset
 
-from transformers import BlipProcessor, BlipForConditionalGeneration
+from transformers import (
+    BlipProcessor, 
+    BlipForConditionalGeneration,
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM
+)
+
 from PIL import Image
 import torch
 
@@ -29,6 +35,10 @@ df = pd.DataFrame(samples)
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
 
+#Load transformer for translating captions from English to Spanish
+trans_tokenizer = AutoTokenizer.from_pretrained("Abhra-loony/english-to-spanish-lang-translation-model")
+trans_model = AutoModelForSeq2SeqLM.from_pretrained("Abhra-loony/english-to-spanish-lang-translation-model")
+
 #Configure captioning function
 def caption_random_image():
 
@@ -42,17 +52,26 @@ def caption_random_image():
     inputs = processor(image, return_tensors="pt")
 
     out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
+    caption_eng = processor.decode(out[0], skip_special_tokens=True)
+
+    # Translate caption from English to Spanish
+    trans_inputs = trans_tokenizer.encode(caption_en, return_tensors="pt")
+    trans_out = trans_model.generate(trans_inputs)
+    caption_es = trans_tokenizer.decode(trans_out[0], skip_special_tokens=True)
 
-    return image, caption
+    return image, caption_eng, caption_es
 
 
 demo = gr.Interface(
     fn=caption_random_image,
     inputs=None,
-    outputs=[gr.Image(type="pil"), gr.Textbox(label="Caption")],
-    title="Random COCO Image Captioning",
-    description="Selects a random COCO image from 20 samples and generates a BLIP caption."
+    outputs=[
+        gr.Image(type="pil", label="Random Image"),
+        gr.Textbox(label="Caption (English)"),
+        gr.Textbox(label="Caption (Spanish)")
+    ],
+    title="Image Captioning (with English to Spanish translation)",
+    description="Selects a random COCO image from 20 samples; generates a BLIP caption; then translates the (English) caption to Spanish."
 )