File size: 3,044 Bytes
e93f267 |
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 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
"""
TinyGPT2 Model Wrapper for easy integration (CPU-friendly)
"""
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
class TinyGPT2Model:
"""
Wrapper for sshleifer/tiny-gpt2 model with caching and optimization
Suitable for CPU-only Hugging Face Spaces
"""
_instance = None
_model = None
_tokenizer = None
def __new__(cls):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def __init__(self):
if TinyGPT2Model._model is None:
self._initialize_model()
def _initialize_model(self):
"""Initialize Tiny-GPT2 model"""
print("Loading TinyGPT2 model...")
model_id = "sshleifer/tiny-gpt2"
# Load tokenizer (no need for token argument, model is public)
TinyGPT2Model._tokenizer = AutoTokenizer.from_pretrained(model_id,token=HUGGINGFACE_TOKEN)
# Load model (no quantization, pure CPU)
TinyGPT2Model._model = AutoModelForCausalLM.from_pretrained(
model_id,token=HUGGINGFACE_TOKEN,
torch_dtype=torch.float32 # Safe for CPU only
)
print("TinyGPT2 model loaded successfully!")
def generate(
self,
prompt: str,
max_length: int = 64,
temperature: float = 0.7,
top_p: float = 0.95
) -> str:
"""Generate response from TinyGPT2"""
# For TinyGPT2, no special prompt formatting needed
formatted_prompt = prompt
# Tokenize
inputs = TinyGPT2Model._tokenizer(
formatted_prompt,
return_tensors="pt",
truncation=True,
max_length=256
)
# Move to CPU (optional, for explicitness)
inputs = {k: v.cpu() for k, v in inputs.items()}
# Generate on CPU
with torch.no_grad():
outputs = TinyGPT2Model._model.generate(
**inputs,
max_new_tokens=max_length,
temperature=temperature,
top_p=top_p,
do_sample=True,
pad_token_id=TinyGPT2Model._tokenizer.eos_token_id
)
# Decode only the newly generated tokens (after the prompt)
response = TinyGPT2Model._tokenizer.decode(
outputs[0][inputs['input_ids'].shape[1]:],
skip_special_tokens=True
)
return response.strip()
def generate_embedding(self, text: str) -> torch.Tensor:
"""Generate embeddings for text using last hidden state"""
inputs = TinyGPT2Model._tokenizer(
text,
return_tensors="pt",
truncation=True,
max_length=256
)
inputs = {k: v.cpu() for k, v in inputs.items()}
with torch.no_grad():
outputs = TinyGPT2Model._model(**inputs, output_hidden_states=True)
embeddings = outputs.hidden_states[-1].mean(dim=1)
return embeddings
|