Update app.py

app.py

@@ -6,21 +6,27 @@ Original file is located at
 """
 
 import os
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import MarianMTModel, MarianTokenizer, GPTNeoForCausalLM, AutoTokenizer
 import gradio as gr
 from PIL import Image, UnidentifiedImageError
 import requests
 import io
 
-# Load translation models
+# Load translation model
 model_name = "Helsinki-NLP/opus-mt-mul-en"
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 model = MarianMTModel.from_pretrained(model_name)
 
-# Define language map
+# Load GPT-Neo model for creative text generation
+gpt_neo_model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
+gpt_neo_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
+
+# Define language map (including new languages)
 language_map = {
     "Tamil": "ta",
-    "Russian": "rus"
+    "Russian": "rus",
+    "Arabic": "ar",
+    "Portuguese": "pt"
 }
 
 def translate_text(input_text, selected_languages):
@@ -53,18 +59,26 @@ def generate_image(prompt):
     else:
         return None
 
+def generate_creative_text(translated_text):
+    prompt = f"Create a creative text based on the following sentence: {translated_text}"
+    inputs = gpt_neo_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=100)
+    output = gpt_neo_model.generate(inputs["input_ids"], max_length=100, do_sample=True, temperature=0.7)
+    creative_text = gpt_neo_tokenizer.decode(output[0], skip_special_tokens=True)
+    return creative_text
+
 def process_input(text_input, selected_languages):
     translated_output = translate_text(text_input, selected_languages)
+    creative_text = generate_creative_text(translated_output)
     image = generate_image(translated_output)
-    return translated_output, image
+    return translated_output, creative_text, image
 
 # Gradio interface
 interface = gr.Interface(
     fn=process_input,
-    inputs=[gr.Textbox(label="Input Text"), gr.CheckboxGroup(choices=["Tamil", "Russian"], label="Select Language")],
-    outputs=[gr.Textbox(label="Translated Text"), gr.Image(label="Generated Image")],
-    title="Multilingual Translation and Image Generation",
-    description="Translate Tamil or Russian text to English and generate an image."
+    inputs=[gr.Textbox(label="Input Text"), gr.CheckboxGroup(choices=["Tamil", "Russian", "Arabic", "Portuguese"], label="Select Language")],
+    outputs=[gr.Textbox(label="Translated Text"), gr.Textbox(label="Creative Text"), gr.Image(label="Generated Image")],
+    title="Multilingual Translation, Creative Text, and Image Generation",
+    description="Translate Tamil, Russian, Arabic, or Portuguese text to English, generate creative text, and generate an image."
 )
 
 interface.launch()