| from typing import Dict, Any |
| import torch |
| from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline |
|
|
| class EndpointHandler: |
| def __init__(self, path=""): |
| |
| self.tokenizer = AutoTokenizer.from_pretrained(path) |
| self.model = AutoModelForCausalLM.from_pretrained(path) |
| |
| |
| self.pipeline = pipeline( |
| "text-generation", |
| model=self.model, |
| tokenizer=self.tokenizer, |
| device=0 if torch.cuda.is_available() else -1 |
| ) |
|
|
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
| prompt_input = data.get("inputs", "") |
| vibe = data.get("vibe", "Open to All Paths") |
| |
| |
| prompt = ( |
| f"#### Human (Vibe: {vibe}): {prompt_input.strip()}\n" |
| f"#### Assistant (Vela - your Camino companion):" |
| ) |
| |
| |
| generation_args = data.get("parameters", {}) |
| generation_args.setdefault("max_new_tokens", 1024) |
| generation_args.setdefault("temperature", 0.2) |
| generation_args.setdefault("top_p", 0.95) |
| generation_args.setdefault("do_sample", True) |
| |
| |
| outputs = self.pipeline( |
| prompt, |
| **generation_args |
| ) |
| |
| return outputs |
