shaheerzk/text_to_sql

1import torch2from transformers import AutoModelForSeq2SeqLM, AutoTokenizer3 4class ModelHandler:5    def __init__(self):6        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")7        self.model = AutoModelForSeq2SeqLM.from_pretrained("shaheerzk/text_to_sql")8        self.tokenizer = AutoTokenizer.from_pretrained("shaheerzk/text_to_sql")9        self.model.to(self.device)10 11    def handle(self, inputs):12        # Preprocess input13        text = inputs.get("text", "")14        inputs = self.tokenizer(text, return_tensors="pt").to(self.device)15 16        # Inference17        with torch.no_grad():18            outputs = self.model.generate(**inputs)19 20        # Post-process output21        generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)22        return {"generated_text": generated_text}23