Coverage for wrapper/slidetranscript/wrapper_slidetranscript.py: 59%
152 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2Slide transcript generation wrapper.
3"""
5import logging
6import json
7import tempfile
8import aiofiles
10from openai import AsyncOpenAI
12from typing import TYPE_CHECKING
13from typing import cast
14from typing import override
15from typing import Dict
16from typing import Optional
17from typing import List
18from typing import Tuple
19from typing import Any
20from typing import AsyncGenerator
22if TYPE_CHECKING:
23 from openai import AsyncStream
24 from openai.types.chat import ChatCompletionChunk
26from wrapper_model import ModelGeneration
28from file_utils import base64_to_binary
29from media_utils import fix_json_like_string
30from file_utils import binary_to_base64
32from transcript_prompts import SYSTEM_PROMPT
34from pptx import Presentation
36from ppt_utils import pptx_to_images
39DEFAULT_MAX_TOKENS = 5000
40DEFAULT_TEMPERATURE = 0.7
43class SlideTranscriptGenerator(ModelGeneration):
44 """A wrapper for slide transcript generation using an LLM."""
46 def __init__(
47 self,
48 llm_url: str = "http://localhost:8000/v1",
49 llm_model: str = "meta-llama/Meta-Llama-3.1-8B",
50 ) -> None:
51 super().__init__("slidetranscript")
52 self.set_llm(llm_model, llm_url)
54 def set_llm(
55 self,
56 llm_model: str,
57 llm_url: str,
58 api_key: str = "n/a",
59 ) -> None:
60 self.llm_model = llm_model
61 self.llm_url = llm_url
62 self.llm_client = AsyncOpenAI(
63 api_key=api_key,
64 base_url=self.llm_url)
65 self.extra_body = dict(guided_decoding_backend="xgrammar")
67 def parse_pptx(
68 self,
69 pptx_path: str
70 ) -> Tuple[Optional[List[str]], Optional[List[str]]]:
71 pptx_texts = []
72 pptx_images: List[str] = []
74 # Add the text from each slide
75 presentation = Presentation(pptx_path)
76 slide_ix = 0
77 for slide in presentation.slides:
78 is_hidden = slide._element.get("show") == "0"
79 if not is_hidden:
80 slide_text = [
81 f"--- SLIDE {slide_ix + 1} ---"
82 ]
83 for shape in slide.shapes:
84 if shape.has_text_frame:
85 for paragraph in shape.text_frame.paragraphs:
86 slide_text.append(paragraph.text)
87 pptx_texts.append("\n".join(slide_text))
88 slide_ix += 1
89 logging.info(f"Detected {slide_ix} slides.")
90 logging.info(f"Extracted text from {len(pptx_texts)} slides.")
92 # Render each of the slides
93 pptx_image_paths = pptx_to_images(
94 pptx_path,
95 output_path="/tmp",
96 )
97 for pptx_image_path in pptx_image_paths:
98 with open(pptx_image_path, "rb") as image_file:
99 image_binary = image_file.read()
100 image_base64 = binary_to_base64(image_binary)
101 pptx_images.append(image_base64)
102 logging.info(f"Rendered {len(pptx_images)} slide images.")
104 return pptx_texts, pptx_images
106 @override
107 async def generate(
108 self,
109 pptx_texts: Optional[List[str]] = None,
110 pptx_images: Optional[List[str]] = None,
111 max_tokens: int = DEFAULT_MAX_TOKENS,
112 temperature: float = DEFAULT_TEMPERATURE,
113 llm_model: Optional[str] = None,
114 llm_url: Optional[str] = None,
115 max_words_per_slide: int = -1,
116 job_id: Optional[str] = None,
117 ) -> List[Dict[str, Any]]:
118 gen_timer = self._new_gen_timer(job_id)
120 self.running = True # We can run in parallel but good to know if we are running
122 if llm_model is not None and llm_url is not None:
123 logging.info(f"Setting LLM model to {llm_model} at {llm_url}.")
124 self.set_llm(llm_model, llm_url)
126 try:
127 ret = []
128 async for slide_text in self.generate_stream(
129 pptx_texts,
130 pptx_images,
131 max_tokens=max_tokens,
132 temperature=temperature,
133 max_words_per_slide=max_words_per_slide,
134 job_id=job_id,
135 ):
136 ret.append(slide_text)
137 return ret
138 finally:
139 gen_timer.end("total")
141 async def generate_stream(
142 self,
143 pptx_texts: Optional[List[str]] = None,
144 pptx_images: Optional[List[str]] = None,
145 max_tokens: int = DEFAULT_MAX_TOKENS,
146 temperature: float = DEFAULT_TEMPERATURE,
147 llm_model: Optional[str] = None,
148 llm_url: Optional[str] = None,
149 max_words_per_slide: int = -1,
150 job_id: Optional[str] = None,
151 ) -> AsyncGenerator[Dict, None]:
152 gen_timer = self._new_gen_timer(job_id)
154 self.running = True # We can run in parallel but good to know if we are running
156 if llm_model is not None and llm_url is not None:
157 logging.info(f"Setting LLM model to {llm_model} at {llm_url}.")
158 self.set_llm(llm_model, llm_url)
160 try:
161 it = 0
162 gen_timer.start("gen_script_stream")
163 gen_timer.start(f"gen_script_stream_{it}")
164 async for slide_text in self.gen_script_stream(
165 pptx_texts,
166 pptx_images,
167 max_words_per_slide=max_words_per_slide,
168 job_id=job_id,
169 ):
170 yield slide_text
171 gen_timer.end(f"gen_script_stream_{it}")
172 it += 1
173 gen_timer.start(f"gen_script_stream_{it}")
174 gen_timer.end(f"gen_script_stream_{it}")
175 gen_timer.end("gen_script_stream")
176 finally:
177 self.running = False
178 gen_timer.end("total")
180 async def gen_script_stream(
181 self,
182 pptx_texts: Optional[List[str]] = None,
183 pptx_images: Optional[List[str]] = None,
184 max_tokens: int = DEFAULT_MAX_TOKENS,
185 temperature: float = DEFAULT_TEMPERATURE,
186 max_words_per_slide: int = -1,
187 job_id: Optional[str] = None,
188 ) -> AsyncGenerator[Dict, None]:
189 if not pptx_texts:
190 raise ValueError("pptx_texts is required to generate slide transcript.")
191 if not pptx_images:
192 raise ValueError("pptx_images is required to generate slide transcript.")
193 if len(pptx_texts) != len(pptx_images):
194 raise ValueError(
195 f"pptx_texts ({len(pptx_texts)}) and "
196 f"pptx_images ({len(pptx_images)}) must have the same length.")
198 contraint_instruction = ""
199 if max_words_per_slide > 0:
200 contraint_instruction += f"Each slide transcript should not exceed {max_words_per_slide} words.\n"
202 # Compose prompt
203 system_prompt_text = (
204 f"{SYSTEM_PROMPT}\n\n"
205 f"{contraint_instruction}"
206 )
208 # Build the messages
209 user_msg: List[Dict[str, Any]] = []
210 for slide_num, (slide_text, slide_image) in enumerate(zip(pptx_texts, pptx_images)):
211 user_msg.append({
212 "type": "text",
213 "text": (
214 f"{slide_text}\n"
215 "The following image corresponds exactly to this slide.\n"
216 )
217 })
218 user_msg.append({
219 "type": "image_url",
220 "image_url": {
221 "url": f"data:image/jpeg;base64,{slide_image}"
222 }
223 })
224 user_msg.append({
225 "type": "text",
226 "text": "Generate the JSONL transcript."
227 })
229 messages: List[Dict[str, Any]] = [
230 {"role": "system", "content": system_prompt_text},
231 {"role": "user", "content": user_msg}
232 ]
234 async with aiofiles.open(f"/tmp/{job_id}_prompt.json", "w") as prompt_file:
235 await prompt_file.write(json.dumps(messages, indent=2))
237 response_stream = cast("AsyncStream[ChatCompletionChunk]", await self.llm_client.chat.completions.create(
238 model=self.llm_model,
239 messages=messages, # type: ignore[arg-type]
240 temperature=temperature,
241 max_tokens=max_tokens,
242 extra_body=self.extra_body,
243 stream=True,
244 ))
246 it = 0
247 buffer_text = ""
248 async for chunk in response_stream:
249 if self.is_interrupted():
250 logging.info("Generation interrupted.")
251 return
253 delta = chunk.choices[0].delta.content or ""
254 buffer_text += delta
255 if delta.endswith("\n"):
256 buffer_text = buffer_text.strip()
257 if buffer_text.startswith("{") and buffer_text.endswith("}"):
258 try:
259 buffer_text = fix_json_like_string(buffer_text)
260 buffer_json = json.loads(buffer_text)
261 yield buffer_json
262 except json.JSONDecodeError as json_error:
263 logging.error(f"JSON error: {json_error} for buffer: {buffer_text}")
264 else:
265 logging.info(f"Ignoring: {buffer_text}")
266 buffer_text = ""
267 it += 1
269 async def warmup(self) -> None:
270 pass # No specific warmup needed for this model
272 async def get_rest_args(
273 self,
274 data_json: Dict[str, str]
275 ) -> Dict[str, Any]:
276 if data_json is None:
277 raise ValueError("Missing JSON body")
279 job_id = data_json.get("job_id", None)
281 pptx_base64 = data_json.get("pptx", None)
282 pptx_texts = None
283 pptx_images = None
284 if pptx_base64 is not None:
285 if not job_id:
286 output_path = tempfile.NamedTemporaryFile(suffix=".pptx", delete=False).name
287 else:
288 output_path = f"/tmp/{job_id}.pptx"
289 pptx_binary = base64_to_binary(pptx_base64)
290 async with aiofiles.open(output_path, "wb") as file:
291 await file.write(pptx_binary)
292 pptx_texts, pptx_images = self.parse_pptx(output_path)
294 rest_args: Dict[str, Any] = {
295 "job_id": job_id,
296 "llm_url": data_json.get("llm_url", "http://localhost:8000/v1"),
297 "llm_model": data_json.get("llm_model", "google/gemma-3-27b-it"),
298 "max_words_per_slide": data_json.get("max_words_per_slide", -1)
299 }
301 if pptx_texts is not None:
302 rest_args["pptx_texts"] = pptx_texts
303 if pptx_images is not None:
304 rest_args["pptx_images"] = pptx_images
306 return {
307 "task": self.model_name,
308 "args": rest_args
309 }