|
|
import gradio as gr |
|
|
import os |
|
|
from src.processing.gemini_processor import GeminiProcessor |
|
|
from src.analysis.coverage_generator import CoverageGenerator |
|
|
from pathlib import Path |
|
|
import logging |
|
|
|
|
|
class ConsoleOutput: |
|
|
def __init__(self): |
|
|
self.messages = [] |
|
|
|
|
|
def write(self, message): |
|
|
self.messages.append(str(message)) |
|
|
return "\n".join(self.messages) |
|
|
|
|
|
def get_output(self): |
|
|
return "\n".join(self.messages) |
|
|
|
|
|
def process_pdf(pdf_file, progress=gr.Progress()): |
|
|
if pdf_file is None: |
|
|
raise gr.Error("Please upload a PDF file") |
|
|
|
|
|
console = ConsoleOutput() |
|
|
|
|
|
|
|
|
logging.basicConfig(level=logging.DEBUG, |
|
|
format='%(asctime)s - %(levelname)s - %(message)s') |
|
|
|
|
|
logger = logging.getLogger() |
|
|
logger.handlers = [] |
|
|
|
|
|
class GradioHandler(logging.Handler): |
|
|
def emit(self, record): |
|
|
msg = self.format(record) |
|
|
console.write(msg) |
|
|
|
|
|
logger.addHandler(GradioHandler()) |
|
|
|
|
|
try: |
|
|
processor = GeminiProcessor() |
|
|
coverage_gen = CoverageGenerator() |
|
|
|
|
|
|
|
|
progress(0.2, desc="Processing screenplay...") |
|
|
cleaned_path = Path("cleaned_screenplay_long.txt") |
|
|
success = processor.process_screenplay(pdf_file.name, str(cleaned_path)) |
|
|
if not success: |
|
|
raise gr.Error("Failed to process screenplay") |
|
|
|
|
|
with open(cleaned_path, 'r') as f: |
|
|
cleaned_text = f.read() |
|
|
|
|
|
|
|
|
progress(0.6, desc="Generating coverage...") |
|
|
success = coverage_gen.generate_coverage(cleaned_path) |
|
|
if not success: |
|
|
raise gr.Error("Failed to generate coverage") |
|
|
|
|
|
with open(Path("coverage.txt"), 'r') as f: |
|
|
coverage = f.read() |
|
|
|
|
|
progress(1.0, desc="Complete!") |
|
|
return cleaned_text, coverage, console.get_output() |
|
|
|
|
|
except Exception as e: |
|
|
error_msg = f"Error: {str(e)}" |
|
|
console.write(error_msg) |
|
|
raise gr.Error(error_msg) |
|
|
|
|
|
with gr.Blocks(title="Screenplay Coverage Generator") as demo: |
|
|
gr.Markdown("# Screenplay Coverage Generator") |
|
|
|
|
|
with gr.Row(): |
|
|
file_input = gr.File(label="Upload Screenplay PDF", file_types=[".pdf"]) |
|
|
|
|
|
with gr.Row(): |
|
|
process_btn = gr.Button("Process Screenplay") |
|
|
|
|
|
with gr.Row(): |
|
|
console = gr.Textbox( |
|
|
label="Console Output", |
|
|
lines=10, |
|
|
max_lines=30, |
|
|
autoscroll=True, |
|
|
container=True, |
|
|
show_copy_button=True |
|
|
) |
|
|
|
|
|
with gr.Tabs(): |
|
|
with gr.TabItem("Cleaned Screenplay"): |
|
|
cleaned_output = gr.Textbox( |
|
|
label="Cleaned Screenplay", |
|
|
lines=10, |
|
|
max_lines=30, |
|
|
show_copy_button=True |
|
|
) |
|
|
with gr.TabItem("Coverage"): |
|
|
coverage_output = gr.Textbox( |
|
|
label="Coverage Document", |
|
|
lines=10, |
|
|
max_lines=30, |
|
|
show_copy_button=True |
|
|
) |
|
|
|
|
|
process_btn.click( |
|
|
fn=process_pdf, |
|
|
inputs=[file_input], |
|
|
outputs=[cleaned_output, coverage_output, console] |
|
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
|
demo.launch() |