sitammeur commited on
Commit
e4b8814
·
verified ·
1 Parent(s): b966f2d

Delete src/app

Browse files
Files changed (3) hide show
  1. src/app/__init__.py +0 -0
  2. src/app/model.py +0 -53
  3. src/app/response.py +0 -77
src/app/__init__.py DELETED
File without changes
src/app/model.py DELETED
@@ -1,53 +0,0 @@
1
- # Necessary imports
2
- import os
3
- import sys
4
- from dotenv import load_dotenv
5
- from typing import Any
6
- import torch
7
- from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
8
-
9
- # Local imports
10
- from src.logger import logging
11
- from src.exception import CustomExceptionHandling
12
-
13
-
14
- # Load the Environment Variables from .env file
15
- load_dotenv()
16
-
17
- # Access token for using the model
18
- access_token = os.environ.get("ACCESS_TOKEN")
19
-
20
-
21
- def load_model_and_processor(model_name: str, device: str) -> Any:
22
- """
23
- Load the model and processor.
24
-
25
- Args:
26
- - model_name (str): The name of the model to load.
27
- - device (str): The device to load the model onto.
28
-
29
- Returns:
30
- - model: The loaded model.
31
- - processor: The loaded processor.
32
- """
33
- try:
34
- # Load the model and processor
35
- model = (
36
- PaliGemmaForConditionalGeneration.from_pretrained(
37
- model_name, torch_dtype=torch.bfloat16, token=access_token
38
- )
39
- .eval()
40
- .to(device)
41
- )
42
- processor = PaliGemmaProcessor.from_pretrained(model_name, token=access_token)
43
-
44
- # Log the successful loading of the model and processor
45
- logging.info("Model and processor loaded successfully.")
46
-
47
- # Return the model and processor
48
- return model, processor
49
-
50
- # Handle exceptions that may occur during model and processor loading
51
- except Exception as e:
52
- # Custom exception handling
53
- raise CustomExceptionHandling(e, sys) from e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
src/app/response.py DELETED
@@ -1,77 +0,0 @@
1
- # Necessary imports
2
- import sys
3
- import PIL.Image
4
- import torch
5
- import gradio as gr
6
- import spaces
7
-
8
- # Local imports
9
- from src.config import device, model_name
10
- from src.app.model import load_model_and_processor
11
- from src.logger import logging
12
- from src.exception import CustomExceptionHandling
13
-
14
-
15
- # Language dictionary
16
- language_dict = {
17
- "English": "en",
18
- "Spanish": "es",
19
- "French": "fr",
20
- }
21
-
22
- # Model and processor
23
- model, processor = load_model_and_processor(model_name, device)
24
-
25
-
26
- @spaces.GPU
27
- def caption_image(
28
- image: PIL.Image.Image, max_new_tokens: int, language: str, sampling: bool
29
- ) -> str:
30
- """
31
- Generates a caption based on the given image using the model.
32
-
33
- Args:
34
- - image (PIL.Image.Image): The input image to be processed.
35
- - max_new_tokens (int): The maximum number of new tokens to generate.
36
- - language (str): The language of the generated caption.
37
- - sampling (bool): Whether to use sampling or not.
38
-
39
- Returns:
40
- str: The generated caption text.
41
- """
42
- try:
43
- # Check if image is None
44
- if not image:
45
- gr.Warning("Please provide an image.")
46
-
47
- # Prepare the inputs
48
- print(language)
49
- language = language_dict[language]
50
- print(language)
51
- prompt = f"<image>caption {language}"
52
- print(prompt)
53
- model_inputs = (
54
- processor(text=prompt, images=image, return_tensors="pt")
55
- .to(torch.bfloat16)
56
- .to(device)
57
- )
58
- input_len = model_inputs["input_ids"].shape[-1]
59
-
60
- # Generate the response
61
- with torch.inference_mode():
62
- generation = model.generate(
63
- **model_inputs, max_new_tokens=max_new_tokens, do_sample=sampling
64
- )
65
- generation = generation[0][input_len:]
66
- decoded = processor.decode(generation, skip_special_tokens=True)
67
-
68
- # Log the successful generation of the caption
69
- logging.info("Caption generated successfully.")
70
-
71
- # Return the generated caption
72
- return decoded
73
-
74
- # Handle exceptions that may occur during caption generation
75
- except Exception as e:
76
- # Custom exception handling
77
- raise CustomExceptionHandling(e, sys) from e