Update app.py

app.py

@@ -1,48 +1,50 @@
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import MarianMTModel, MarianTokenizer, GPT2LMHeadModel, GPT2Tokenizer
 
 def translate(text, target_language):
-    language_codes = {
-        "Spanish": "es",
-        "French (European)": "fr",
-        "French (Canadian)": "fr",
-        "Italian": "it",
-        "Ukrainian": "uk",
-        "Portuguese (Brazilian)": "pt_BR",
-        "Portuguese (European)": "pt",
-        "Russian": "ru",
-        "Chinese": "zh",
-        "Dutch": "nl",
-        "German": "de",
-        "Arabic": "ar",
-        "Hebrew": "he",
-        "Greek": "el"
-    }
-    target_language_code = language_codes[target_language]
-    model_name = f'helsinki-nlp/opus-mt-en-{target_language_code}'
-    tokenizer = MarianTokenizer.from_pretrained(model_name)
-    model = MarianMTModel.from_pretrained(model_name)
-    inputs = tokenizer(text, return_tensors="pt")
-    outputs = model.generate(**inputs)
-    translation = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return translation
+    # ... (keep the existing code for translation here)
+
+def generate_text(prompt):
+    model_name = 'gpt2'
+    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+    model = GPT2LMHeadModel.from_pretrained(model_name)
+    inputs = tokenizer.encode(prompt, return_tensors='pt')
+    outputs = model.generate(inputs, max_length=100, num_return_sequences=1)
+    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return generated_text
 
 language_options = [
-    "Spanish", "French (European)", "French (Canadian)", "Italian", "Ukrainian",
-    "Portuguese (Brazilian)", "Portuguese (European)", "Russian", "Chinese",
-    "Dutch", "German", "Arabic", "Hebrew", "Greek"
+    # ... (keep the existing language options here)
 ]
 
-iface = gr.Interface(
+iface_translation = gr.Interface(
     fn=translate,
+    # ... (keep the existing translation inputs and outputs here)
+)
+
+iface_generation = gr.Interface(
+    fn=generate_text,
+    inputs=gr.inputs.Textbox(lines=5, label="Enter a prompt for text generation:"),
+    outputs=gr.outputs.Textbox(label="Generated Text"),
+)
+
+# Combine the two interfaces into a single Gradio interface
+iface_combined = gr.Interface(
+    [translate, generate_text],
     inputs=[
-        gr.inputs.Textbox(lines=5, label="Enter text to translate:"),
+        gr.inputs.Textbox(lines=5, label="Enter text to translate / generate:", default="Enter text to translate here."),
         gr.inputs.Dropdown(choices=language_options, label="Target Language"),
     ],
-    outputs=gr.outputs.Textbox(label="Translated Text"),
+    outputs=[
+        gr.outputs.Textbox(label="Translated Text / Generated Text"),
+    ],
+    title="Translation and Text Generation",
+    description="Choose a target language to translate English text or leave it as 'None' for text generation.",
+    examples=[["Translate this text to French.", "French (European)"]]
 )
 
-iface.launch()
+iface_combined.launch()
+