add language detection, improve the ability to follow the reference copy

app.py

@@ -153,6 +153,65 @@ def detect_features(image_paths, garment_type, language="English"):
         return "", []
     
 
+import re
+def parse_structure(struct_ref):
+    languages = ["_"] * (len(struct_ref) // 2)
+    types = ["_"] * (len(struct_ref) // 2)
+    for si in range(0, len(struct_ref), 2):
+        parts = re.findall('[a-zA-Z\n ]+', struct_ref[si])
+        for idx, part in enumerate(parts):
+            if "language" in part.lower():
+                lang = parts[idx + 1].strip()
+                languages[si // 2] = lang
+            if "type" in part.lower():
+                type = parts[idx + 1].strip()
+                types[si // 2] = type
+
+    return types, languages
+
+
+def detect_language(texts, model):
+    langs = ["_"] * len(texts)
+
+    try:
+        messages = []
+        lang_map = {}
+        for i, text in enumerate(texts):
+            if len(text.strip()) > 0:
+                lang_mess = [HumanMessage(content=f"What is the language of the following text? Output the language only. "
+                                        f"\n ```{text}```")]
+                print(f"{lang_mess=}")
+                messages.append(lang_mess)
+                lang_map[i] = len(messages) - 1 
+        detected_langs = model.batch(messages)
+        print(f"{detected_langs=}")
+        for k, v in lang_map.items():
+            langs[k] = detected_langs[v].content
+    except Exception as e:
+        print(e.__class__, e)
+        traceback.print_exc()
+
+    return langs
+
+
+def get_language(struct_lang, copy_lang):
+    if struct_lang != "_":
+        return struct_lang
+    if copy_lang != "_":
+        return copy_lang
+    return "English"
+
+
+def get_model(model_name):
+    if model_name.startswith("gpt"):
+        chat = ChatOpenAI(model=model_name)
+    elif model_name.startswith("claude"):
+        chat = ChatAnthropic(model_name=model_name, anthropic_api_key=os.environ["ANTHROPIC_API_KEY"])
+    else:
+        chat = ChatGroq(model_name=model_name, api_key=os.environ["GROQ_API_KEY"])
+    return chat
+    
+
 def generate(*data):
     global visible
     print("visible", visible)
@@ -171,6 +230,17 @@ def generate(*data):
     print(f"{glossary=}")
     print(f"{struct_ref=}")
 
+    chat = get_model(model)
+
+    types, struct_languages = parse_structure(struct_ref)
+    copy_languages = detect_language([struct_ref[2 * i + 1] for i in range(visible + 1)], model=chat) 
+    languages = [get_language(struct_lang=struct_lang, copy_lang=copy_lang) for struct_lang, copy_lang in zip(struct_languages, copy_languages)]
+
+    print("Struct languages--------------------------------------------\n", struct_languages)
+    print("Copy languages--------------------------------------------\n", copy_languages)
+    print("Languages--------------------------------------------\n", languages)
+    # print("Types--------------------------------------------", types)
+
     image_features, base64_images = detect_features(image, garment_type)
     detected_features = ""
     intended_use = ""
@@ -180,12 +250,6 @@ def generate(*data):
         detected_features = ", ".join(image_features["features"])
         intended_use = "Intended use: " + ", ".join(image_features["intended_use"])
         print(f"Detected features: {detected_features}, Intended use: {intended_use}, Alt text: {alt_texts}")
-    if model.startswith("gpt"):
-        chat = ChatOpenAI(model=model)
-    elif model.startswith("claude"):
-        chat = ChatAnthropic(model_name=model, anthropic_api_key=os.environ["ANTHROPIC_API_KEY"])
-    else:
-        chat = ChatGroq(model_name=model, api_key=os.environ["GROQ_API_KEY"])
 
     batch = []
     for i in range(visible + 1):
@@ -195,18 +259,19 @@ def generate(*data):
             # visible = i 
             # break
             messages = [
-                SystemMessage(content=f"""You are a helpful assistant that writes about products for ecommerce websites."""),
+                SystemMessage(content=f"""You are a helpful assistant that writes about products for ecommerce websites. Make sure to write in {languages[i]} language."""),
                 HumanMessage(content=f"""Write a product description with the following features.
                                     Make sure that the description follows the structure of the reference structure. 
                                     Make sure to use markdown format for the output.
                                     Make sure that the entire output is written entirely in language defined in the reference structure. 
+                                    Make sure to output the product description only, do not include any preceeding text like "Here is your product description".
                                     Use language that is suitable for the type of document specified in the reference structure. 
                                     Use a consistent tone of voice throughout the text. 
                                     If the reference text is not empty, write the product description in the tone of voice of the reference text. 
-                                    Make sure to output the product description only, do not include any preceeding text like "Here is your product description".
                                     Do not include any part of the reference structure in the output.
                                     Do not use any of the excluded words in the output.
                                     Make sure to include all of the included words in the output.
+                                    Do not hallucinate any information.
                                     
                                     {feature + detected_features} 
                                     {intended_use} 
@@ -221,6 +286,38 @@ def generate(*data):
     response = chat.batch(batch, temperature=temperature)
     print(response)
 
+    batch = []
+    rewrite_map = {}
+    for i in range(visible + 1):
+        structure = struct_ref[2 * i]
+        reference = struct_ref[2 * i + 1]  
+        if len(reference.strip()) > 0:
+            messages = [
+                SystemMessage(content=f"""You are a helpful assistant that writes about products for ecommerce websites. You write in {languages[i]} language."""),
+                HumanMessage(content=f"""Rewrite the following product description in the style and tone of voice
+                              of the reference product description.
+                             Make sure that the structure and length of the output is similar to the reference product description.
+                             Make sure that the output is written in {languages[i]} language.
+                             Output the product description in markdown format. 
+
+                             Product description to rewirte: 
+                             ```{response[i].content}```
+
+                             Reference product description:
+                             ```{reference}```
+                             """)
+            ]
+            batch.append(messages)  
+            rewrite_map[i] = len(batch) - 1
+
+    print("Rewrite_map", rewrite_map)
+    print("Rewriting")
+    re_response = chat.batch(batch, temperature=temperature)
+    for i in range(len(re_response)):
+        print(f"Original: {response[i].content}")
+        print(f"Rewritten: {re_response[i].content}")
+    response = [re_response[rewrite_map[i]] if i in rewrite_map else response[i] for i in range(visible + 1)]
+    print("Done rewriting")
 
     description = "\n---\n".join([msg.content for msg in response])
     md_content = description