File size: 4,486 Bytes
3b01827 | 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 101 102 103 104 105 106 107 108 109 110 111 112 | from typing import Dict, List, Any
from transformers import pipeline
import torch
from datetime import datetime
import pytz
from pathlib import Path
class EndpointHandler:
def __init__(self, model_path: str = "", tokenizer_path: str = None):
"""
Initialize the endpoint handler with model and system instruction
Args:
model_path: Path to the model
tokenizer_path: Path to the tokenizer (if different from model_path)
"""
# Set default tokenizer path if not provided
if tokenizer_path is None:
tokenizer_path = model_path
# Initialize the pipeline
self.pipeline = pipeline(
"text-generation",
model=model_path,
tokenizer=tokenizer_path,
max_length=2000,
device_map="cuda" if torch.cuda.is_available() else "cpu"
)
# System instruction
self.system_instruction =
"""You are an AI that generates a user-friendly user interface based on a user query using the µUI script language. You must output the µUI script covering the sufficient interface that is needed without any other content.
You must include the information mentioned by the user usable to construct your interface.
Only when the request is unclear or require mandatory context or information to generate UI for it, you generate the right UI to collect such missing information or clarify the context. You also must use µUI (MicroUI) Language for the information collection UI generation as well.
Don't output any other content except µUI script. """
# Default location
self.default_location = "San Francisco, USA"
def get_formatted_time(self, timezone: str = "America/Los_Angeles") -> str:
"""Get formatted time for given timezone"""
try:
tz = pytz.timezone(timezone)
current_time = datetime.now(tz)
return current_time.strftime('%Y-%m-%d %H:%M:%S %Z')
except:
# Default to PT if timezone is invalid
tz = pytz.timezone("America/Los_Angeles")
current_time = datetime.now(tz)
return current_time.strftime('%Y-%m-%d %H:%M:%S %Z')
def format_message(self, query: str, location: str = None, time: str = None) -> str:
"""Format the user message with location and time"""
location = location if location else self.default_location
time = time if time else self.get_formatted_time()
return f"Current User query: {query}\nCurrent Location: {location}\nCurrent Time: {time}"
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
"""
Handle the inference request
Args:
data: Dictionary containing:
- inputs (str): User query
- location (str, optional): User location
- time (str, optional): Custom time
Returns:
List[Dict]: Model response with metadata
"""
try:
# Extract inputs
if isinstance(data, dict):
inputs = data.pop("inputs", None)
if inputs is None:
# If no inputs field, treat entire data as input
inputs = data
location = data.pop("location", None)
time = data.pop("time", None)
else:
# If data is not a dict, treat it as the input
inputs = data
location = None
time = None
# Format messages
messages = [
{"role": "system", "content": self.system_instruction},
{"role": "user", "content": self.format_message(inputs, location, time)}
]
# Run inference
output = self.pipeline(messages)
# Extract and format response
response = {
"generated_text": output[0]["generated_text"][2]["content"],
"metadata": {
"location": location if location else self.default_location,
"time": time if time else self.get_formatted_time()
}
}
return [response]
except Exception as e:
return [{"error": f"Inference failed: {str(e)}"}]
|