File size: 4,189 Bytes
2575aee
 
 
c353670
2575aee
 
93992d2
2575aee
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
import fitz  # PyMuPDF
import gradio as gr
import json
from langchain_text_splitters import RecursiveCharacterTextSplitter
from groq import Groq
import os

# Load API key
API_KEY = os.getenv('Groq')
if not API_KEY:
    raise ValueError("API Key is missing! Set the environment variable 'GROQ_API_KEY'.")

# Initialize Groq Client
client = Groq(api_key=API_KEY)

# Prompt Template
PROMPT_TEMPLATE = """
You are an expert screenplay analyst. Convert the following film script text into the JSON structure below:
{json_structure}

Script Text:
{text}

Provide only the JSON response.
""".strip()

# Define the JSON structure to be extracted
JSON_STRUCTURE = {
    "scenes": [
        {
            "scene_heading": "",
            "location": "",
            "time_of_day": "",
            "characters": [],
            "emotions": [],
            "summary": "",
            "dialogues": [
                {
                    "character": "",
                    "dialogue_text": "",
                    "tone": ""
                }
            ]
        }
    ],
    "overall_emotional_arc": [],
    "story_beats": {
        "setup": "",
        "inciting_incident": "",
        "climax": "",
        "resolution": ""
    }
}

# Function to extract text from PDF
def extract_text_from_pdf(pdf_file):
    text = ""
    try:
        with open(pdf_file.name, 'rb') as f:
            doc = fitz.open(stream=f.read(), filetype="pdf")
            for page in doc:
                text += page.get_text() + "\n"
    except Exception as e:
        return f"Error reading PDF: {e}"
    return text.strip()

# Function to split text into chunks
def split_text_into_chunks(text, chunk_size=2000):
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=100)
    return splitter.split_text(text)

# Function to call Groq API
def call_llm_api(text):
    prompt = PROMPT_TEMPLATE.format(json_structure=json.dumps(JSON_STRUCTURE, indent=2), text=text)
    
    try:
        response = client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama-3-3-70b-vision",  # You can also test llama-3-70b-versatile
        )

        raw_content = response.choices[0].message.content.strip()
        
        # Clean JSON formatting
        if raw_content.startswith("```json") and raw_content.endswith("```"):
            raw_content = raw_content[7:-3].strip()
        
        return json.loads(raw_content)
    except Exception as e:
        return {"error": f"API call failed: {e}"}

# Function to merge JSON chunks
def merge_json_chunks(chunks):
    combined_result = JSON_STRUCTURE.copy()
    combined_result["scenes"] = []
    combined_result["overall_emotional_arc"] = []
    
    for chunk in chunks:
        result = call_llm_api(chunk)
        if "error" in result:
            continue
        
        # Merge scenes
        if "scenes" in result:
            combined_result["scenes"].extend(result["scenes"])
        
        # Merge emotional arc
        if "overall_emotional_arc" in result:
            combined_result["overall_emotional_arc"].extend(result["overall_emotional_arc"])
        
        # Merge story beats only once (first time we encounter valid values)
        for beat in combined_result["story_beats"].keys():
            if result.get("story_beats", {}).get(beat) and not combined_result["story_beats"][beat]:
                combined_result["story_beats"][beat] = result["story_beats"][beat]
    
    return combined_result

# Gradio interface function
def gradio_interface(file):
    pdf_text = extract_text_from_pdf(file)
    if pdf_text.startswith("Error"):
        return {"error": pdf_text}
    
    chunks = split_text_into_chunks(pdf_text)
    extracted_data = merge_json_chunks(chunks)
    return extracted_data

# Gradio UI
iface = gr.Interface(
    fn=gradio_interface,
    inputs=gr.File(label="Upload Film Script PDF"),
    outputs="json",
    title="ScriptWhisper - Screenplay Structure & Emotion Extractor",
    description="Upload a screenplay PDF to extract scene structure, emotional arc, and story beats."
)

# Launch the app
if __name__ == "__main__":
    iface.launch()