Upload model.py with huggingface_hub

a381a76 verified over 1 year ago

4 kB

import torch
from transformers import PreTrainedModel, PretrainedConfig
from PIL import Image
import numpy as np
from typing import Optional, Union, List
import json

class Pix2TextConfig(PretrainedConfig):
    model_type = "pix2text"
    
    def __init__(
        self,
        vocab_size=30000,
        hidden_size=768,
        num_attention_heads=12,
        num_hidden_layers=12,
        intermediate_size=3072,
        max_position_embeddings=512,
        dropout_prob=0.1,
        **kwargs
    ):
        super().__init__(**kwargs)
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_attention_heads = num_attention_heads
        self.num_hidden_layers = num_hidden_layers
        self.intermediate_size = intermediate_size
        self.max_position_embeddings = max_position_embeddings
        self.dropout_prob = dropout_prob

class Pix2TextModel(PreTrainedModel):
    config_class = Pix2TextConfig
    
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        
        # Basit bir CNN encoder
        self.image_encoder = torch.nn.Sequential(
            torch.nn.Conv2d(3, 64, 3, padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2),
            torch.nn.Conv2d(64, 128, 3, padding=1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(2),
            torch.nn.Conv2d(128, 256, 3, padding=1),
            torch.nn.ReLU(),
            torch.nn.AdaptiveAvgPool2d((8, 8)),
            torch.nn.Flatten(),
            torch.nn.Linear(256 * 8 * 8, config.hidden_size)
        )
        
        # Text decoder
        self.text_decoder = torch.nn.Sequential(
            torch.nn.Linear(config.hidden_size, config.intermediate_size),
            torch.nn.ReLU(),
            torch.nn.Dropout(config.dropout_prob),
            torch.nn.Linear(config.intermediate_size, config.vocab_size)
        )
        
        # Basit tokenizer için vocab
        self.vocab = {str(i): i for i in range(10)}  # 0-9 rakamları
        self.vocab.update({chr(i): i+10 for i in range(ord('a'), ord('z')+1)})  # a-z
        self.vocab.update({chr(i): i+36 for i in range(ord('A'), ord('Z')+1)})  # A-Z
        self.vocab.update({' ': 62, '.': 63, ',': 64, '!': 65, '?': 66})
        self.inv_vocab = {v: k for k, v in self.vocab.items()}
        
    def forward(self, pixel_values, labels=None):
        # Görüntüyü encode et
        image_features = self.image_encoder(pixel_values)
        
        # Text'e decode et
        logits = self.text_decoder(image_features)
        
        loss = None
        if labels is not None:
            loss_fct = torch.nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
            
        return {"loss": loss, "logits": logits}
    
    def predict(self, image: Union[Image.Image, np.ndarray]) -> str:
        """Görüntüden metin çıkarma fonksiyonu"""
        if isinstance(image, Image.Image):
            image = np.array(image)
            
        # Görüntüyü ön işle
        if len(image.shape) == 3:
            image = torch.tensor(image).permute(2, 0, 1).float() / 255.0
        else:
            image = torch.tensor(image).unsqueeze(0).float() / 255.0
            
        image = image.unsqueeze(0)  # Batch dimension
        
        with torch.no_grad():
            outputs = self.forward(image)
            logits = outputs["logits"]
            
            # En yüksek olasılıklı tokenları seç
            predicted_ids = torch.argmax(logits, dim=-1)
            
            # Token'ları metne çevir
            text = ""
            for token_id in predicted_ids[0][:10]:  # İlk 10 token
                if token_id.item() in self.inv_vocab:
                    text += self.inv_vocab[token_id.item()]
                    
        return text

	import torch
	from transformers import PreTrainedModel, PretrainedConfig
	from PIL import Image
	import numpy as np
	from typing import Optional, Union, List
	import json

	class Pix2TextConfig(PretrainedConfig):
	model_type = "pix2text"

	def __init__(
	self,
	vocab_size=30000,
	hidden_size=768,
	num_attention_heads=12,
	num_hidden_layers=12,
	intermediate_size=3072,
	max_position_embeddings=512,
	dropout_prob=0.1,
	**kwargs
	):
	super().__init__(**kwargs)
	self.vocab_size = vocab_size
	self.hidden_size = hidden_size
	self.num_attention_heads = num_attention_heads
	self.num_hidden_layers = num_hidden_layers
	self.intermediate_size = intermediate_size
	self.max_position_embeddings = max_position_embeddings
	self.dropout_prob = dropout_prob

	class Pix2TextModel(PreTrainedModel):
	config_class = Pix2TextConfig

	def __init__(self, config):
	super().__init__(config)
	self.config = config

	# Basit bir CNN encoder
	self.image_encoder = torch.nn.Sequential(
	torch.nn.Conv2d(3, 64, 3, padding=1),
	torch.nn.ReLU(),
	torch.nn.MaxPool2d(2),
	torch.nn.Conv2d(64, 128, 3, padding=1),
	torch.nn.ReLU(),
	torch.nn.MaxPool2d(2),
	torch.nn.Conv2d(128, 256, 3, padding=1),
	torch.nn.ReLU(),
	torch.nn.AdaptiveAvgPool2d((8, 8)),
	torch.nn.Flatten(),
	torch.nn.Linear(256 * 8 * 8, config.hidden_size)
	)

	# Text decoder
	self.text_decoder = torch.nn.Sequential(
	torch.nn.Linear(config.hidden_size, config.intermediate_size),
	torch.nn.ReLU(),
	torch.nn.Dropout(config.dropout_prob),
	torch.nn.Linear(config.intermediate_size, config.vocab_size)
	)

	# Basit tokenizer için vocab
	self.vocab = {str(i): i for i in range(10)} # 0-9 rakamları
	self.vocab.update({chr(i): i+10 for i in range(ord('a'), ord('z')+1)}) # a-z
	self.vocab.update({chr(i): i+36 for i in range(ord('A'), ord('Z')+1)}) # A-Z
	self.vocab.update({' ': 62, '.': 63, ',': 64, '!': 65, '?': 66})
	self.inv_vocab = {v: k for k, v in self.vocab.items()}

	def forward(self, pixel_values, labels=None):
	# Görüntüyü encode et
	image_features = self.image_encoder(pixel_values)

	# Text'e decode et
	logits = self.text_decoder(image_features)

	loss = None
	if labels is not None:
	loss_fct = torch.nn.CrossEntropyLoss()
	loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))

	return {"loss": loss, "logits": logits}

	def predict(self, image: Union[Image.Image, np.ndarray]) -> str:
	"""Görüntüden metin çıkarma fonksiyonu"""
	if isinstance(image, Image.Image):
	image = np.array(image)

	# Görüntüyü ön işle
	if len(image.shape) == 3:
	image = torch.tensor(image).permute(2, 0, 1).float() / 255.0
	else:
	image = torch.tensor(image).unsqueeze(0).float() / 255.0

	image = image.unsqueeze(0) # Batch dimension

	with torch.no_grad():
	outputs = self.forward(image)
	logits = outputs["logits"]

	# En yüksek olasılıklı tokenları seç
	predicted_ids = torch.argmax(logits, dim=-1)

	# Token'ları metne çevir
	text = ""
	for token_id in predicted_ids[0][:10]: # İlk 10 token
	if token_id.item() in self.inv_vocab:
	text += self.inv_vocab[token_id.item()]

	return text