Spaces:
Sleeping
Sleeping
File size: 6,534 Bytes
910dbfd d2d66c1 910dbfd d2d66c1 c144d94 910dbfd c144d94 d2d66c1 910dbfd c144d94 910dbfd c144d94 910dbfd df844fd c144d94 910dbfd c144d94 910dbfd c144d94 d2d66c1 c144d94 910dbfd c144d94 d2d66c1 c144d94 d2d66c1 c144d94 d2d66c1 c144d94 910dbfd c144d94 910dbfd c144d94 910dbfd c144d94 910dbfd c144d94 |
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 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 |
from abc import ABC, abstractmethod
from typing import Type, TypeVar
import base64
import os
import json
from doc2json import process_docx
import fitz
from PIL import Image
import io
import boto3
from botocore.config import Config
import re
# constants
log_to_console = False
use_document_message_type = False # AWS document message type usage
LLMClass = TypeVar('LLMClass', bound='LLM')
class LLM:
@staticmethod
def create_llm(model: str) -> Type[LLMClass]:
return LLM()
def generate_body(self, message, history):
messages = []
# AWS API requires strict user, assi, user, ... sequence
lastTypeHuman = False
for human, assi in history:
if human:
if lastTypeHuman:
last_msg = messages.pop()
user_msg_parts = last_msg["content"]
else:
user_msg_parts = []
if isinstance(human, tuple):
user_msg_parts.extend(self._process_file(human[0]))
else:
user_msg_parts.extend([{"text": human}])
messages.append({"role": "user", "content": user_msg_parts})
lastTypeHuman = True
if assi:
messages.append({"role": "assistant", "content": [{"text": assi}]})
lastTypeHuman = False
user_msg_parts = []
if message.text:
user_msg_parts.append({"text": message.text})
if message.files:
for file in message.files:
user_msg_parts.extend(self._process_file(file.path))
if user_msg_parts:
messages.append({"role": "user", "content": user_msg_parts})
return messages
def _process_file(self, file_path):
if use_document_message_type and self._is_supported_document_type(file_path):
return [self._create_document_message(file_path)]
else:
return self._encode_file(file_path)
def _is_supported_document_type(self, file_path):
supported_extensions = ['.pdf', '.csv', '.doc', '.docx', '.xls', '.xlsx', '.html', '.txt', '.md']
return os.path.splitext(file_path)[1].lower() in supported_extensions
def _create_document_message(self, file_path):
with open(file_path, 'rb') as file:
file_content = file.read()
file_name = re.sub(r'[^a-zA-Z0-9\s\-\(\)\[\]]', '', os.path.basename(file_path))[:200].strip() or "unnamed_file"
file_extension = os.path.splitext(file_path)[1][1:] # Remove the dot
return {
"document": {
"name": file_name,
"format": file_extension,
"source": {
"bytes": file_content
}
}
}
def _encode_file(self, fn: str) -> list:
if fn.endswith(".docx"):
return [{"text": process_docx(fn)}]
elif fn.endswith(".pdf"):
return self._process_pdf_img(fn)
else:
with open(fn, mode="rb") as f:
content = f.read()
if isinstance(content, bytes):
try:
# try to add as image
image_data = self._encode_image(content)
return [{"image": image_data}]
except:
# not an image, try text
content = content.decode('utf-8', 'replace')
else:
content = str(content)
fname = os.path.basename(fn)
return [{"text": f"``` {fname}\n{content}\n```"}]
def _process_pdf_img(self, pdf_fn: str):
pdf = fitz.open(pdf_fn)
message_parts = []
for page in pdf.pages():
# Create a transformation matrix for rendering at the calculated scale
mat = fitz.Matrix(0.6, 0.6)
# Render the page to a pixmap
pix = page.get_pixmap(matrix=mat, alpha=False)
# Convert pixmap to PIL Image
img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
# Convert PIL Image to bytes
img_byte_arr = io.BytesIO()
img.save(img_byte_arr, format='PNG')
img_byte_arr = img_byte_arr.getvalue()
# Append the message parts
message_parts.append({"text": f"Page {page.number} of file '{pdf_fn}'"})
message_parts.append({"image": {
"format": "png",
"source": {"bytes": img_byte_arr}
}})
pdf.close()
return message_parts
def _encode_image(self, image_data):
# Get the first few bytes of the image data.
magic_number = image_data[:4]
# Check the magic number to determine the image type.
if magic_number.startswith(b'\x89PNG'):
image_type = 'png'
elif magic_number.startswith(b'\xFF\xD8'):
image_type = 'jpeg'
elif magic_number.startswith(b'GIF89a'):
image_type = 'gif'
elif magic_number.startswith(b'RIFF'):
if image_data[8:12] == b'WEBP':
image_type = 'webp'
else:
# Unknown image type.
raise Exception("Unknown image type")
else:
# Unknown image type.
raise Exception("Unknown image type")
return {
"format": image_type,
"source": {"bytes": image_data}
}
def read_response(self, response_stream):
for event in response_stream:
if 'contentBlockDelta' in event:
yield event['contentBlockDelta']['delta']['text']
if 'messageStop' in event:
if log_to_console:
print(f"\nStop reason: {event['messageStop']['stopReason']}")
if 'metadata' in event:
metadata = event['metadata']
if 'usage' in metadata and log_to_console:
print("\nToken usage:")
print(f"Input tokens: {metadata['usage']['inputTokens']}")
print(f"Output tokens: {metadata['usage']['outputTokens']}")
print(f"Total tokens: {metadata['usage']['totalTokens']}") |