Spaces:
Sleeping
Sleeping
| from src.config.settings import GEMINI_API_KEY, CHATBOT_NAME, MODEL_ID, MODEL_TEMPERATURE, MODEL_OPTIONS, IMAGE_WIDTH, IMAGE_HEIGHT, SYSTEM_INSTRUCTION | |
| from typing import Dict, List, Optional | |
| from PIL import Image | |
| from google import genai | |
| from google.genai import types | |
| import time | |
| def preprocess_stop_sequences(stop_sequences: str) -> Optional[List[str]]: | |
| return [seq.strip() for seq in stop_sequences.split(",")] if stop_sequences else None | |
| def preprocess_image(image: Image.Image) -> Image.Image: | |
| image_height = int(image.height * IMAGE_WIDTH / image.width) | |
| return image.resize((IMAGE_WIDTH, image_height)) | |
| def user(text_prompt: str, chatbot: List[Dict[str, str]]): | |
| return "", chatbot + [{"role": "user", "content": text_prompt}] | |
| def bot( | |
| model_name: str, | |
| image_prompt: Optional[Image.Image], | |
| temperature: float, | |
| max_output_tokens: int, | |
| stop_sequences: str, | |
| top_k: int, | |
| top_p: float, | |
| chatbot: List[Dict[str, str]] | |
| ): | |
| if not GEMINI_API_KEY: | |
| chatbot.append({"role": "assistant", "content": "GEMINI_API_KEY is not set. Please add it to your .env file."}) | |
| yield chatbot | |
| return | |
| client = genai.Client(api_key=GEMINI_API_KEY) | |
| # Gradio v6 may store content as a list of parts or a plain string | |
| raw_content = chatbot[-1].get("content") if chatbot else None | |
| if isinstance(raw_content, list): | |
| text_prompt = " ".join( | |
| part.get("text", "") if isinstance(part, dict) else str(part) | |
| for part in raw_content | |
| ).strip() or None | |
| elif isinstance(raw_content, str): | |
| text_prompt = raw_content.strip() or None | |
| else: | |
| text_prompt = None | |
| if not text_prompt and not image_prompt: | |
| chatbot.append({"role": "assistant", "content": "Prompt cannot be empty. Please provide input text or an image."}) | |
| yield chatbot | |
| return | |
| elif image_prompt and not text_prompt: | |
| text_prompt = "Describe the image" | |
| elif image_prompt and text_prompt: | |
| text_prompt = f"{text_prompt}. Also, analyze the provided image." | |
| contents = [text_prompt] if image_prompt is None else [text_prompt, preprocess_image(image_prompt)] | |
| config = types.GenerateContentConfig( | |
| system_instruction=SYSTEM_INSTRUCTION, | |
| temperature=temperature, | |
| max_output_tokens=max_output_tokens, | |
| stop_sequences=preprocess_stop_sequences(stop_sequences), | |
| top_k=top_k, | |
| top_p=top_p, | |
| safety_settings=[ | |
| types.SafetySetting( | |
| category="HARM_CATEGORY_DANGEROUS_CONTENT", | |
| threshold="BLOCK_ONLY_HIGH", | |
| ) | |
| ], | |
| ) | |
| chatbot.append({"role": "assistant", "content": ""}) | |
| try: | |
| for chunk in client.models.generate_content_stream( | |
| model=model_name, | |
| contents=contents, | |
| config=config, | |
| ): | |
| if chunk.text: | |
| for i in range(0, len(chunk.text), 10): | |
| chatbot[-1]["content"] += chunk.text[i:i + 10] | |
| time.sleep(0.01) | |
| yield chatbot | |
| except Exception as e: | |
| chatbot[-1]["content"] = f"Error occurred: {str(e)}" | |
| yield chatbot | |
| return | |