import os
import datetime
import requests
import re
import pandas as pd
import gradio as gr
import threading
import uuid
import queue
import time
import fitz # PyMuPDF for reading PDF files
from transformers import AutoTokenizer
from mistralai import Mistral
from huggingface_hub import InferenceClient
# ------------------------------
# Helper functions and globals
# ------------------------------
sheet_data = None
file_name = None
def debug_print(message: str):
print(f"[{datetime.datetime.now().isoformat()}] {message}", flush=True)
def initialize_tokenizer():
try:
return AutoTokenizer.from_pretrained("gpt2")
except Exception as e:
debug_print("Failed to initialize tokenizer: " + str(e))
return None
global_tokenizer = initialize_tokenizer()
def count_tokens(text: str) -> int:
if global_tokenizer:
try:
return len(global_tokenizer.encode(text))
except Exception:
return len(text.split())
return len(text.split())
def generate_response(prompt: str, model_name: str, sheet_data: str) -> str:
full_prompt = f"{prompt}\n\nSheet Data:\n{sheet_data}" # Append loaded text to prompt
if "Mistral" in model_name:
mistral_api_key = os.getenv("MISTRAL_API_KEY")
if not mistral_api_key:
raise ValueError("MISTRAL_API_KEY environment variable not set.")
mistral_client = Mistral(api_key=mistral_api_key)
response = mistral_client.chat.complete(
model="mistral-small-latest",
messages=[{"role": "user", "content": full_prompt}],
temperature=0.7,
top_p=0.95
)
return response.choices[0].message.content
elif "Meta-Llama" in model_name:
hf_api_token = os.getenv("HF_API_TOKEN")
if not hf_api_token:
raise ValueError("HF_API_TOKEN environment variable not set.")
client = InferenceClient(token=hf_api_token)
response = client.text_generation(
full_prompt,
model="meta-llama/Meta-Llama-3-8B-Instruct",
temperature=0.7,
top_p=0.95,
max_new_tokens=512
)
return response
else:
raise ValueError("Invalid model selection. Please choose either 'Mistral-API' or 'Meta-Llama-3'.")
def process_query(prompt: str, model_name: str):
global sheet_data
if sheet_data is None:
sheet_data = get_sheet_data()
full_prompt = f"{prompt}\n\nSheet Data:\n{sheet_data}"
debug_print(f"Processing query with model {model_name}: {full_prompt}")
response = generate_response(prompt, model_name, sheet_data)
input_tokens = count_tokens(prompt + "\n\n" + sheet_data)
output_tokens = count_tokens(response)
return response, f"Input tokens: {input_tokens}", f"Output tokens: {output_tokens}"
def ui_process_query(prompt, model_name):
return process_query(prompt, model_name)
# ------------------------------
# Cleaning Functions
# ------------------------------
def clean_text(text: str, remove_spaces: bool, remove_headers_footers: bool, lowercase: bool, remove_special: bool) -> str:
"""
Cleans the given text based on the provided options.
"""
# Remove extra spaces & newlines
if remove_spaces:
text = re.sub(r'\s+', ' ', text).strip()
# Remove headers/footers: a simple heuristic to remove lines that repeat
if remove_headers_footers:
lines = text.split('\n')
freq = {}
for line in lines:
line_stripped = line.strip()
if line_stripped:
freq[line] = freq.get(line, 0) + 1
lines = [line for line in lines if freq.get(line, 0) <= 1]
text = "\n".join(lines)
if lowercase:
text = text.lower()
if remove_special:
text = re.sub(r'[^a-zA-Z0-9\s]', '', text)
return text
def execute_cleaning(text: str, remove_spaces: bool, remove_headers: bool, lowercase: bool, remove_special: bool) -> str:
if not text or text.strip() == "":
return "No text available for cleaning."
cleaned = clean_text(text, remove_spaces, remove_headers, lowercase, remove_special)
return cleaned
# ------------------------------
# Global variables for background jobs
# ------------------------------
jobs = {}
results_queue = queue.Queue()
last_job_id = None
# ------------------------------
# Job management functions
# ------------------------------
def get_job_list():
job_list_md = "### π Submitted Jobs\n\n"
if not jobs:
return "No jobs found. Submit a query or load files to create jobs."
sorted_jobs = sorted(
[(job_id, job_info) for job_id, job_info in jobs.items()],
key=lambda x: x[1].get("start_time", 0),
reverse=True
)
for job_id, job_info in sorted_jobs:
status = job_info.get("status", "unknown")
job_type = job_info.get("type", "unknown")
query = job_info.get("query", "")
start_time = job_info.get("start_time", 0)
time_str = datetime.datetime.fromtimestamp(start_time).strftime("%Y-%m-%d %H:%M:%S")
query_preview = query[:30] + "..." if query and len(query) > 30 else query or "N/A"
if status == "processing":
status_formatted = f"β³ {status}"
elif status == "completed":
status_formatted = f"β
{status}"
else:
status_formatted = f"β {status}"
if job_type == "query":
job_list_md += f"- [{job_id}](javascript:void) - {time_str} - {status_formatted} - Query: {query_preview}\n"
else:
job_list_md += f"- [{job_id}](javascript:void) - {time_str} - {status_formatted} - File Load Job\n"
return job_list_md
def get_sheet_data():
global sheet_data
return sheet_data if sheet_data else "No data loaded."
def process_in_background(job_id, func, args):
result = func(*args)
results_queue.put((job_id, result))
debug_print(f"Job {job_id} finished processing in background.")
def submit_query_async(query, model_choice=None):
global last_job_id
global sheet_data
if not query:
return ("Please enter a non-empty query", "", "Input tokens: 0", "Output tokens: 0", "", "", get_job_list())
job_id = str(uuid.uuid4())
debug_print(f"Starting async job {job_id} for query: {query}")
threading.Thread(
target=process_in_background,
args=(job_id, process_query, [query, model_choice or "Mistral-API"])
).start()
jobs[job_id] = {
"status": "processing",
"type": "query",
"start_time": time.time(),
"query": query,
"model": model_choice or "Mistral-API"
}
last_job_id = job_id
return (
f"π Query submitted and processing in the background (Job ID: {job_id}).\n\n"
f"Use the 'Check Job Status' section to view results.",
f"Job ID: {job_id}",
f"Input tokens: {count_tokens(query)}",
"Output tokens: pending",
job_id,
query,
get_job_list()
)
def job_selected(job_id):
if job_id in jobs:
return job_id, jobs[job_id].get("query", "No query for this job")
return job_id, "Job not found"
def refresh_job_list():
return get_job_list()
def sync_model_dropdown(value):
return value
def check_job_status(job_id):
if not job_id:
html_response = ""
return html_response, "", "", "", ""
try:
while not results_queue.empty():
completed_id, result = results_queue.get_nowait()
if completed_id in jobs:
jobs[completed_id]["status"] = "completed"
jobs[completed_id]["result"] = result
jobs[completed_id]["end_time"] = time.time()
debug_print(f"Job {completed_id} completed and stored in jobs dictionary")
except queue.Empty:
pass
if job_id not in jobs:
html_response = "Job not found. Please check the ID and try again.
"
f"
β³ Query is still being processed (elapsed: {elapsed_time:.1f}s). Please check again shortly.
"
f"
"
f"
β
Response: {result[0]}
"
f"
Processing time: {processing_time:.1f}s
"
f"
Job status: {job['status']}