Update app.py

app.py

@@ -10,20 +10,17 @@ model_name = "Helsinki-NLP/opus-mt-mul-en"
 tokenizer = MarianTokenizer.from_pretrained(model_name)
 model = MarianMTModel.from_pretrained(model_name)
 
-# Load GPT-Neo model for creative text generation
+# Load GPT-Neo model for creative content 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)
+# Define language map
 language_map = {
     "Tamil": "ta",
-    "Russian": "rus",
-    "Arabic": "ar",
-    "Portuguese": "pt"
+    "Russian": "rus"
 }
 
 def translate_text(input_text, selected_languages):
-    """Translate input text into English based on the selected language."""
     if not selected_languages:
         return "Please select at least one language."
     
@@ -32,16 +29,13 @@ def translate_text(input_text, selected_languages):
     lang_prefix = f">>{lang_code}<< "
     text_with_lang = lang_prefix + input_text
     inputs = tokenizer(text_with_lang, return_tensors="pt", padding=True)
-    
-    # Generate translated tokens
     translated_tokens = model.generate(**inputs)
     translation = tokenizer.decode(translated_tokens[0], skip_special_tokens=True)
     return translation
 
 def generate_image(prompt):
-    """Generate an image based on the provided prompt."""
     API_URL = "https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev"
-    hf_token = os.getenv("HF_TOKEN")  # Ensure to set this environment variable
+    hf_token = os.getenv("HF_TOKEN")
     headers = {"Authorization": f"Bearer {hf_token}"}
     
     response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
@@ -57,17 +51,13 @@ def generate_image(prompt):
         return None
 
 def generate_creative_text(translated_text):
-    """Generate creative text based on the translated sentence."""
-    prompt = f"Create a creative text based on the following sentence: {translated_text}"
+    prompt = f"Create a creative story based on the following sentence: {translated_text}"
     inputs = gpt_neo_tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=100)
-    
-    # Generate creative text
-    output = gpt_neo_model.generate(inputs["input_ids"], max_length=100, do_sample=True, temperature=0.7)
+    output = gpt_neo_model.generate(inputs["input_ids"], max_length=150, 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):
-    """Process the input text: translate, generate creative text, and generate an image."""
     translated_output = translate_text(text_input, selected_languages)
     creative_text = generate_creative_text(translated_output)
     image = generate_image(translated_output)
@@ -77,16 +67,16 @@ def process_input(text_input, selected_languages):
 interface = gr.Interface(
     fn=process_input,
     inputs=[
-        gr.Textbox(label="Input Text"), 
-        gr.CheckboxGroup(choices=["Tamil", "Russian", "Arabic", "Portuguese"], label="Select Language")
+        gr.Textbox(label="Input Text"),
+        gr.CheckboxGroup(choices=["Tamil", "Russian"], label="Select Language")
     ],
     outputs=[
-        gr.Textbox(label="Translated Text"), 
-        gr.Textbox(label="Creative Text"), 
+        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."
+    title="Multilingual Translation and Image Generation",
+    description="Translate Tamil or Russian text to English, generate creative content, and create an image."
 )
 
 interface.launch()