Spaces:
Build error
Build error
| import gradio as gr | |
| import pandas as pd | |
| from openai import OpenAI | |
| import os | |
| import pdfplumber | |
| import pytesseract | |
| from PIL import Image | |
| # Set your OpenAI API key using environment variables for security | |
| OPENAI_API_KEY = os.getenv('OPEN_API_KEY') | |
| client = OpenAI( | |
| api_key=OPENAI_API_KEY | |
| # os.getenv('OPENAI_API_KEY'), | |
| ) | |
| # Function to extract text from PDFs | |
| def extract_text_from_pdf(file_path): | |
| text = "" | |
| try: | |
| with pdfplumber.open(file_path) as pdf: | |
| for page in pdf.pages: | |
| text += page.extract_text() + "\n" | |
| except Exception as e: | |
| print(f"Error extracting text from PDF: {str(e)}") | |
| return text | |
| # Function to extract text from images using OCR | |
| def extract_text_from_image(image): | |
| try: | |
| text = pytesseract.image_to_string(image) | |
| return text | |
| except Exception as e: | |
| print(f"Error extracting text from image: {str(e)}") | |
| return "" | |
| # Function to infer table and column names from extracted text | |
| def infer_table_columns_from_text(text): | |
| messages = [ | |
| {"role": "system", "content": "You are a helpful assistant that can interpret text to identify table and column names."}, | |
| {"role": "user", "content": f"The following text is extracted from a document: '{text}'. Identify any table names and corresponding column names in the text."} | |
| ] | |
| response = client.chat.completions.create( | |
| model="gpt-4", | |
| messages=messages, | |
| max_tokens=300, | |
| temperature=0.5 | |
| ) | |
| return response.choices[0].message.content.strip() | |
| # Function to classify columns using GPT and pre-defined classifications | |
| def classify_columns_from_text(text, classification_file='classifications.xlsx'): | |
| # Load classification rules | |
| classification_df = pd.read_excel(classification_file) | |
| # print("Classification DataFrame Columns:\n", classification_df.columns) | |
| # Infer table and column names using GPT-4 | |
| inferred_structure = infer_table_columns_from_text(text) | |
| # print("GPT Inferred Structure:\n", inferred_structure) | |
| classified_columns = [] | |
| inferred_lines = inferred_structure.split("\n") | |
| table_name = None | |
| for line in inferred_lines: | |
| line = line.strip() | |
| if "typically be found in a" in line and "table" in line: | |
| table_name = line.split("typically be found in a ")[-1].split(" table")[0].strip() | |
| print(f"Inferred Table Name: {table_name}") | |
| elif line and line[0].isdigit(): | |
| column_name = line.split(".")[1].strip() | |
| print(f"Checking for Column: '{column_name}' in Table: '{table_name}'") | |
| # Find the closest match using GPT-4 to compare the inferred column name with the classification file's columns | |
| best_match = None | |
| best_similarity = 0 | |
| for col in classification_df['Column names'].unique(): | |
| prompt_messages = [ | |
| {"role": "system", "content": "You are a helpful assistant that can compare and identify similarity between two strings."}, | |
| {"role": "user", "content": f"On a scale of 0 to 1, how similar are these two column names? Column 1: '{column_name}', Column 2: '{col}'. Respond with only the similarity score."} | |
| ] | |
| response = client.chat.completions.create( | |
| model="gpt-4", | |
| messages=prompt_messages, | |
| max_tokens=10, | |
| temperature=0 | |
| ) | |
| raw_response = response.choices[0].message.content.strip() | |
| # print("Raw GPT-4 Response:", raw_response) | |
| try: | |
| similarity_score = float(response.choices[0].message.content.strip()) | |
| except ValueError: | |
| print("Error: Unable to convert GPT-4 response to a float. Response was:", response.choices[0].message['content'].strip()) | |
| similarity_score = 0.0 # Default to 0 if conversion fails | |
| if similarity_score > best_similarity: | |
| best_similarity = similarity_score | |
| best_match = col | |
| # print(f"Best match for '{column_name}' is '{best_match}' with similarity {best_similarity}") | |
| if best_match: | |
| classification_match = classification_df[classification_df['Column names'].str.strip().str.lower() == best_match.strip().lower()]['Classification'] | |
| if not classification_match.empty: | |
| classification = classification_match.values[0] | |
| print(f"Match Found: '{column_name}' classified as '{classification}'") | |
| else: | |
| classification = "General Use" | |
| print(f"No Match Found for Column: '{column_name}' - Defaulting to 'General Use'") | |
| else: | |
| classification = "General Use" | |
| classified_columns.append({ | |
| 'Table Name': table_name, | |
| 'Column Name': column_name, | |
| 'Classification': classification | |
| }) | |
| result_df = pd.DataFrame(classified_columns) | |
| print("Resulting Classified Columns DataFrame:\n", result_df) | |
| return result_df | |
| def process_file(file): | |
| text = "" | |
| file_extension = file.name.split(".")[-1].lower() | |
| if file_extension in ["pdf"]: | |
| # Extract text from PDF | |
| text = extract_text_from_pdf(file) | |
| elif file_extension in ["jpg", "jpeg", "png", "tiff"]: | |
| # Extract text from image | |
| image = Image.open(file) | |
| text = extract_text_from_image(image) | |
| else: | |
| # Assume it's an Excel file | |
| input_df = pd.read_excel(file) | |
| text = "\n".join(input_df.apply(lambda row: " ".join(row.values.astype(str)), axis=1)) | |
| # Classify columns based on the extracted text | |
| result_df = classify_columns_from_text(text) | |
| return result_df | |
| # Create a Gradio interface | |
| iface = gr.Interface( | |
| fn=process_file, | |
| inputs=gr.File(label="Upload Excel, PDF, or Image File"), | |
| outputs=gr.Dataframe(label="Classified Data"), | |
| title="Data Classification using GPT and Pre-defined Rules", | |
| description="Upload an Excel, PDF, or Image file to classify its columns based on predefined rules and GPT-4." | |
| ) | |
| # Launch the Gradio app with a public link | |
| iface.launch(share=True) | |