Update app.py

app.py

@@ -63,17 +63,18 @@ def generate_image(prompt):
         print(f"An error occurred: {e}")
         return None
 
-# Load GPT-Neo model for creative text generation
+# Import necessary libraries for Mistral model
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-gpt_neo_tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neo-125M")
-gpt_neo_model = AutoModelForCausalLM.from_pretrained("EleutherAI/gpt-neo-125M")
+# Load Mistral model and tokenizer
+mistral_tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
+mistral_model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
 
-# Function to generate creative text based on translated text
+# Function to generate creative text based on translated text using Mistral
 def generate_creative_text(translated_text):
-    input_ids = gpt_neo_tokenizer(translated_text, return_tensors='pt').input_ids
-    generated_text_ids = gpt_neo_model.generate(input_ids, max_length=100)
-    creative_text = gpt_neo_tokenizer.decode(generated_text_ids[0], skip_special_tokens=True)
+    input_ids = mistral_tokenizer(translated_text, return_tensors='pt').input_ids
+    generated_text_ids = mistral_model.generate(input_ids, max_length=100)
+    creative_text = mistral_tokenizer.decode(generated_text_ids[0], skip_special_tokens=True)
     return creative_text
 
 # Function to handle the full workflow