File size: 2,715 Bytes
070c576 aa834c5 070c576 aa834c5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 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)), gr.update(choices=list(df.columns)), gr.update(choices=list(df.columns)), gr.update(choices=list(df.columns)), df.head(5)
else:
return gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[]), gr.update(choices=[]), pd.DataFrame() |