mirror of
https://github.com/run-llama/llama-hub.git
synced 2025-08-14 11:41:56 +00:00
e97bb81915
* cr * cr * cr --------- Co-authored-by: Jerry Liu <jerry@robustintelligence.com> Co-authored-by: Jesse Zhang <jessetanzhang@gmail.com>
104 lines
3.4 KiB
Python
104 lines
3.4 KiB
Python
"""Read Microsoft PowerPoint files."""
|
|
|
|
import os
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
from llama_index.readers.base import BaseReader
|
|
from llama_index.readers.schema.base import Document
|
|
|
|
|
|
class PptxReader(BaseReader):
|
|
"""Powerpoint reader.
|
|
|
|
Extract text, caption images, and specify slides.
|
|
|
|
"""
|
|
|
|
def __init__(self, caption_images: Optional[bool] = False) -> None:
|
|
"""Init reader."""
|
|
self.caption_images = caption_images
|
|
if caption_images:
|
|
from transformers import (AutoTokenizer, VisionEncoderDecoderModel,
|
|
ViTFeatureExtractor)
|
|
|
|
model = VisionEncoderDecoderModel.from_pretrained(
|
|
"nlpconnect/vit-gpt2-image-captioning"
|
|
)
|
|
feature_extractor = ViTFeatureExtractor.from_pretrained(
|
|
"nlpconnect/vit-gpt2-image-captioning"
|
|
)
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
"nlpconnect/vit-gpt2-image-captioning"
|
|
)
|
|
|
|
self.parser_config = {
|
|
"feature_extractor": feature_extractor,
|
|
"model": model,
|
|
"tokenizer": tokenizer,
|
|
}
|
|
|
|
def generate_image_caption(self, tmp_image_file: str) -> str:
|
|
"""Generate text caption of image."""
|
|
if not self.caption_images:
|
|
return ""
|
|
|
|
import torch
|
|
from PIL import Image
|
|
|
|
model = self.parser_config["model"]
|
|
feature_extractor = self.parser_config["feature_extractor"]
|
|
tokenizer = self.parser_config["tokenizer"]
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
model.to(device)
|
|
|
|
max_length = 16
|
|
num_beams = 4
|
|
gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
|
|
|
|
i_image = Image.open(tmp_image_file)
|
|
if i_image.mode != "RGB":
|
|
i_image = i_image.convert(mode="RGB")
|
|
|
|
pixel_values = feature_extractor(
|
|
images=[i_image], return_tensors="pt"
|
|
).pixel_values
|
|
pixel_values = pixel_values.to(device)
|
|
|
|
output_ids = model.generate(pixel_values, **gen_kwargs)
|
|
|
|
preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
|
return preds[0].strip()
|
|
|
|
def load_data(
|
|
self,
|
|
file: Path,
|
|
extra_info: Optional[Dict] = None,
|
|
) -> List[Document]:
|
|
"""Parse file."""
|
|
from pptx import Presentation
|
|
|
|
presentation = Presentation(file)
|
|
result = ""
|
|
for i, slide in enumerate(presentation.slides):
|
|
result += f"\n\nSlide #{i}: \n"
|
|
for shape in slide.shapes:
|
|
if self.caption_images and hasattr(shape, "image"):
|
|
image = shape.image
|
|
# get image "file" contents
|
|
image_bytes = image.blob
|
|
# temporarily save the image to feed into model
|
|
image_filename = f"tmp_image.{image.ext}"
|
|
with open(image_filename, "wb") as f:
|
|
f.write(image_bytes)
|
|
result += (
|
|
f"\n Image: {self.generate_image_caption(image_filename)}\n\n"
|
|
)
|
|
|
|
os.remove(image_filename)
|
|
if hasattr(shape, "text"):
|
|
result += f"{shape.text}\n"
|
|
|
|
return [Document(result, extra_info=extra_info)]
|