Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,705 Bytes
2cfa8e4 |
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 |
from typing import List, Union, Optional
from indexify_extractor_sdk import Content, Extractor, Feature
from pydantic import BaseModel, Field
import os
import google.generativeai as genai
from pdf2image import convert_from_path
import tempfile
import mimetypes
class GeminiExtractorConfig(BaseModel):
model_name: Optional[str] = Field(default='gemini-1.5-flash-latest')
key: Optional[str] = Field(default=None)
prompt: str = Field(default='You are a helpful assistant.')
query: Optional[str] = Field(default=None)
class GeminiExtractor(Extractor):
name = "tensorlake/gemini"
description = "An extractor that let's you use LLMs from Gemini."
system_dependencies = []
input_mime_types = ["text/plain", "application/pdf", "image/jpeg", "image/png"]
def __init__(self):
super(GeminiExtractor, self).__init__()
def extract(self, content: Content, params: GeminiExtractorConfig) -> List[Union[Feature, Content]]:
contents = []
model_name = params.model_name
key = params.key
prompt = params.prompt
query = params.query
if content.content_type in ["application/pdf"]:
with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
temp_file.write(content.data)
file_path = temp_file.name
images = convert_from_path(file_path)
image_files = []
for i in range(len(images)):
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_image_file:
images[i].save(temp_image_file.name, 'JPEG')
image_files.append(temp_image_file.name)
elif content.content_type in ["image/jpeg", "image/png"]:
image_files = []
suffix = mimetypes.guess_extension(content.content_type)
with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_image_file:
temp_image_file.write(content.data)
file_path = temp_image_file.name
image_files.append(file_path)
else:
text = content.data.decode("utf-8")
if query is None:
query = text
file_path = None
def upload_to_gemini(path, mime_type=None):
file = genai.upload_file(path, mime_type=mime_type)
print(f"Uploaded file '{file.display_name}' as: {file.uri}")
return file
if ('GEMINI_API_KEY' not in os.environ) and (key is None):
response_content = "The GEMINI_API_KEY environment variable is not present."
else:
if ('GEMINI_API_KEY' in os.environ) and (key is None):
genai.configure(api_key=os.environ["GEMINI_API_KEY"])
else:
genai.configure(api_key=key)
generation_config = { "temperature": 1, "top_p": 0.95, "top_k": 64, "max_output_tokens": 8192, "response_mime_type": "text/plain", }
model = genai.GenerativeModel( model_name=model_name, generation_config=generation_config, )
if file_path:
files = [upload_to_gemini(image_file, mime_type="image/jpeg") for image_file in image_files]
chat_session = model.start_chat( history=[ { "role": "user", "parts": files, }, ] )
response = chat_session.send_message(prompt)
else:
chat_session = model.start_chat( history=[ ] )
response = chat_session.send_message(prompt + " " + query)
response_content = response.text
contents.append(Content.from_text(response_content))
return contents |