Update app.py

app.py

@@ -1,22 +1,39 @@
 import gradio as gr
-from transformers import MarianMTModel, MarianTokenizer
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load the translation model and tokenizer
-model_name = 'Helsinki-NLP/opus-mt-en-es'
-model = MarianMTModel.from_pretrained(model_name)
-tokenizer = MarianTokenizer.from_pretrained(model_name)
+# Load the Mistral model and tokenizer
+model_name = 'mistral/mistral-7b'  # Replace with your specific Mistral model
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-def translate_text(text):
-    inputs = tokenizer.encode(text, return_tensors="pt")
-    translated = model.generate(inputs, max_length=50)
+# Define a function to handle translation
+def translate_text(text, src_lang, tgt_lang):
+    # Prepare the input text for the model
+    # For simplicity, we'll just append source and target languages to the input text
+    input_text = f"{src_lang} to {tgt_lang}: {text}"
+    inputs = tokenizer(input_text, return_tensors="pt")
+    translated = model.generate(**inputs)
     translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
     return translated_text
 
-def chat(message):
-    translated_message = translate_text(message)
+def chat(message, src_lang, tgt_lang):
+    translated_message = translate_text(message, src_lang, tgt_lang)
     return translated_message
 
+# Define the language options
+languages = ["English", "French", "German", "Spanish", "Chinese"]  # Extend this list as needed
+
 # Create the Gradio interface
-interface = gr.Interface(fn=chat, inputs="text", outputs="text")
+interface = gr.Interface(
+    fn=chat,
+    inputs=[
+        gr.inputs.Textbox(label="Enter text"),
+        gr.inputs.Dropdown(choices=languages, label="Source Language"),
+        gr.inputs.Dropdown(choices=languages, label="Target Language")
+    ],
+    outputs="text",
+    title="Mistral Translation Chatbot",
+    description="Translate text between different languages using the Mistral model."
+)
 
 interface.launch()