Buckets:
| from typing import Optional, Union, Dict, Any, Tuple, List | |
| import torch | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| class BaseModelManager: | |
| def __init__(self, clamp_device_to_cpu: bool = False): | |
| self.base_model_id = "gpt2" | |
| self.device = ( | |
| "cpu" | |
| if clamp_device_to_cpu | |
| else ("cuda" if torch.cuda.is_available() else "cpu") | |
| ) | |
| self.tokenizer = self.setup_tokenizer() | |
| self.base_model = self.load_model() | |
| def setup_tokenizer(self): | |
| tkzr = AutoTokenizer.from_pretrained(self.base_model_id) | |
| if tkzr.pad_token is None: | |
| tkzr.pad_token = tkzr.eos_token | |
| return tkzr | |
| def load_model(self): | |
| model = AutoModelForCausalLM.from_pretrained( | |
| self.base_model_id, | |
| torch_dtype=torch.float32, | |
| device_map="auto" if self.device != "cpu" else None, | |
| torch_compile=False, | |
| ) | |
| if self.device == "cpu": | |
| model = model.to(self.device) | |
| return model | |
| def standard_inference( | |
| self, | |
| model, | |
| tokenizer, | |
| prompt: str, | |
| max_new_tokens: int = 15, | |
| temperature: float = 0.7, | |
| ) -> str: | |
| inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=128) | |
| inputs = {k: v.to(self.device) for k, v in inputs.items()} | |
| with torch.no_grad(): | |
| outputs = model.generate( | |
| **inputs, | |
| max_new_tokens=max_new_tokens, | |
| temperature=temperature, | |
| pad_token_id=tokenizer.pad_token_id, | |
| do_sample=True, | |
| top_p=0.9 | |
| ) | |
| full_response = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
| generated_part = full_response[len(prompt) :].strip() | |
| answer = generated_part.split("\n")[0].split(" ")[0].strip() | |
| return answer | |
Xet Storage Details
- Size:
- 1.93 kB
- Xet hash:
- 27a845fc82bb7a1ca840e466feb95aa886a90d9261249b46a2b6a95bd36560b1
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.