Spaces:
Running
Running
# Gemini Functions | |
import gradio as gr | |
import os | |
import google.generativeai as genai | |
import json | |
def initialize_gemini(): | |
try: | |
genai.configure(api_key=os.getenv("GEMINI_API_KEY")) | |
generation_config = { | |
"temperature": 1, | |
"top_p": 0.95, | |
"top_k": 40, | |
"max_output_tokens": 8192, | |
"response_mime_type": "text/plain", | |
} | |
model = genai.GenerativeModel( | |
model_name="gemini-1.5-pro", | |
generation_config=generation_config, | |
) | |
return model | |
except Exception as e: | |
raise gr.Error(f"Error initializing Gemini: {str(e)}") | |
def create_prompt(extracted_text: str, path_to_data_to_extract: str) -> str: | |
# load data to extract | |
with open(path_to_data_to_extract, 'r', encoding='utf-8') as file: | |
data_to_extract = json.load(file) | |
prompt = f"""Tu es un assistant juridique expert en analyse de documents judiciaires français. | |
Je vais te fournir le contenu d'un document judiciaire extrait d'un PDF. | |
Ta tâche est d'analyser ce texte et d'en extraire les informations suivantes de manière précise : | |
{json.dumps(data_to_extract, indent=2, ensure_ascii=False)} | |
Voici quelques règles à suivre : | |
- Si une information n'est pas présente dans le texte, indique "Non spécifié" pour cette catégorie. | |
- Pour les noms des parties (demandeurs et défendeurs, et leurs avocats), liste tous ceux que tu trouves | |
- Assure-toi de différencier correctement les demandeurs des défendeurs. | |
- Si tu n'es pas sûr d'une information, indique-le clairement. | |
Présente tes résultats sous forme de JSON, en utilisant les catégories mentionnées ci-dessus. | |
Voici le contenu du document : | |
{extracted_text.strip()} | |
Analyse ce texte et fournis-moi les informations demandées au format JSON uniquement.""".strip() | |
return prompt | |
def extract_data_with_gemini(text_file_path: str, path_to_data_to_extract: str) -> dict: | |
try: | |
# Initialize Gemini | |
model = initialize_gemini() | |
# Read the extracted text | |
with open(text_file_path, 'r', encoding='utf-8') as f: | |
extracted_text = f.read() | |
# Create prompt and get response | |
prompt = create_prompt(extracted_text, path_to_data_to_extract) | |
response = model.generate_content(prompt) | |
# Parse the JSON response | |
try: | |
# Extract JSON from the response text | |
json_str = response.text | |
if "json" in json_str.lower(): | |
json_str = json_str.split("json")[1].split("```")[0] | |
elif "```" in json_str: | |
json_str = json_str.split("```")[1] | |
result = json.loads(json_str) | |
except: | |
result = {"error": "Failed to parse JSON response", "raw_response": response.text} | |
return result | |
except Exception as e: | |
raise gr.Error(f"Error in Gemini processing: {str(e)}") | |