improve text generation

app.py

@@ -44,7 +44,6 @@ def get_text_gen_pipeline():
                    model="mistralai/Mistral-7B-Instruct-v0.3",
                    max_new_tokens=2048,
                    device=device,
-                   #model_kwargs={"add_prefix_space": False}
               )
          except Exception as e:
               print(f"Error loading text generation model: {e}")
@@ -58,11 +57,16 @@ def refine_prompt(prompt):
         return "Text generation model is unavailable."
     try:
         messages = [
-            {"role": "system", "content": "You are a product designer. You will get a basic prompt of product request and you need to imagine a new product design to satisfy that need. Produce an extended description of product front view that will be used by a text-to-image AI (Flux) to generate a visual. The prompt should be in the form of a product description, not a story, maximum 2048 tokens. Mention the visual aesthetics. Background should be a white background."},
+            {"role": "system", "content": "You are a product designer specialized in text-to-image generation. You will get a basic prompt of product request and you need to imagine a new product design to satisfy that need. Generate a visual product description that will then be used by a text-to-image AI (Flux) to suggest a visual. The prompt should be in the form of a strict product description, not a story, maximum 2048 tokens. You need to mention explicitly the visual aesthetics (ex: photo realistic, high quality, etc). Background should be a full white background."},           {"role": "system", "content": "You are a product designer specialized in text-to-image generation. You will get a basic prompt of product request and you need to imagine a new product design to satisfy that need. Genrerate a product description of product front view that will be used by a text-to-image AI (Flux) to generate a visual. The prompt should be in the form of a product description, not a story, maximum 2048 tokens. You need to mention explicitly the visual aesthetics (ex: photo realistic, high quality, etc). Background should be a full white background."},
             {"role": "user", "content": prompt},
         ]
         refined_prompt = text_gen(messages)
-        return refined_prompt
+        # Extract just the assistant's content from the response
+        try:
+            assistant_content = refined_prompt[0]['generated_text'][-1]['content']
+            return assistant_content
+        except (KeyError, IndexError):
+            return "Error: Unexpected response format from the model"
     except Exception as e:
         return f"Error refining prompt: {str(e)}"