pengjunjie/ragflow
1
0
Fork 1
mirror of https://github.com/infiniflow/ragflow.git synced 2025-12-08 20:42:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

Refactor:Improve BytesIO usage for GeminiCV (#10042)

Browse Source 
### What problem does this PR solve?
Improve BytesIO usage for GeminiCV

### Type of change
- [x] Refactoring
This commit is contained in:
Stephen Hu
2025-09-11 11:07:15 +08:00
committed by GitHub
parent 8a09f07186
commit 1936ad82d2
1 changed files with 12 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
rag/llm/cv_model.py
View File

@ -521,24 +521,24 @@ class GeminiCV(Base):
else "Please describe the content of this picture, like where, when, who, what happen. If it has number data, please extract them out." else "Please describe the content of this picture, like where, when, who, what happen. If it has number data, please extract them out."
) )
b64 = self.image2base64(image) b64 = self.image2base64(image)
img = open(BytesIO(base64.b64decode(b64))) with BytesIO(base64.b64decode(b64)) as bio:
input = [prompt, img] img = open(bio)
res = self.model.generate_content(input) input = [prompt, img]
img.close() res = self.model.generate_content(input)
return res.text, res.usage_metadata.total_token_count img.close()
return res.text, res.usage_metadata.total_token_count
def describe_with_prompt(self, image, prompt=None): def describe_with_prompt(self, image, prompt=None):
from PIL.Image import open from PIL.Image import open
b64 = self.image2base64(image) b64 = self.image2base64(image)
vision_prompt = prompt if prompt else vision_llm_describe_prompt() vision_prompt = prompt if prompt else vision_llm_describe_prompt()
img = open(BytesIO(base64.b64decode(b64))) with BytesIO(base64.b64decode(b64)) as bio:
input = [vision_prompt, img] img = open(bio)
res = self.model.generate_content( input = [vision_prompt, img]
input, res = self.model.generate_content(input)
) img.close()
img.close() return res.text, res.usage_metadata.total_token_count
return res.text, res.usage_metadata.total_token_count
def chat(self, system, history, gen_conf, images=[]): def chat(self, system, history, gen_conf, images=[]):
generation_config = dict(temperature=gen_conf.get("temperature", 0.3), top_p=gen_conf.get("top_p", 0.7)) generation_config = dict(temperature=gen_conf.get("temperature", 0.3), top_p=gen_conf.get("top_p", 0.7))