Create excel_chat.py

excel_chat.py

@@ -0,0 +1,58 @@
+import gradio as gr
+from mistralai.client import MistralClient
+from mistralai.models.chat_completion import ChatMessage
+import os
+import pandas as pd
+import numpy as np
+
+def chat_with_mistral(source_cols, dest_col, prompt, tdoc_name, excel_file, url):
+
+    df = pd.read_excel(excel_file)
+    api_key = os.environ["MISTRAL_API_KEY"]
+    model = "mistral-small"  # Use "Mistral-7B-v0.2" for "mistral-tiny"
+
+
+    client = MistralClient(api_key=api_key)
+
+    source_columns = source_cols#.split(", ")  # Split input into multiple variables
+    df[dest_col] = ""
+    try:
+      file_name = url.split("/")[-2] + ".xlsx"
+    except:
+      file_name = excel_file
+
+    if tdoc_name != '':
+        filtered_df = df[df['File'] == tdoc_name]
+        if not filtered_df.empty:
+            concatenated_content = "\n\n".join(f"{column_name}: {filtered_df[column_name].iloc[0]}" for column_name in source_columns)
+            messages = [ChatMessage(role="user", content=f"Using the following content: {concatenated_content}"), ChatMessage(role="user", content=prompt)]
+            chat_response = client.chat(model=model, messages=messages)
+            filtered_df.loc[filtered_df.index[0], dest_col] = chat_response.choices[0].message.content
+            # Update the DataFrame with the modified row
+            df.update(filtered_df)
+            # Write the updated DataFrame to the Excel file
+            df.to_excel(file_name, index=False)
+            return file_name, df.head(5)
+        else:
+            return file_name, df.head(5)
+    else:
+        for index, row in df.iterrows():
+            concatenated_content = "\n\n".join(f"{column_name}: {row[column_name]}" for column_name in source_columns)
+            # Check if the concatenated content is not empty
+            print('test')
+            if not concatenated_content == "\n\n".join(f"{column_name}: nan" for column_name in source_columns):
+                print('c bon')
+                messages = [ChatMessage(role="user", content=f"Using the following content: {concatenated_content}"), ChatMessage(role="user", content=prompt)]
+                chat_response = client.chat(model=model, messages=messages)
+                df.at[index, dest_col] = chat_response.choices[0].message.content
+
+        df.to_excel(file_name, index=False)
+        return file_name, df.head(5)
+
+
+def get_columns(file):
+    if file is not None:
+        df = pd.read_excel(file)
+        return gr.update(choices=list(df.columns)), df.head(5)
+    else:
+        return gr.update(choices=[]), pd.DataFrame()