shaheerzk commited on
Commit
971be23
·
verified ·
1 Parent(s): b0182f4

Upload handler.py

Browse files
Files changed (1) hide show
  1. handler.py +22 -0
handler.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
3
+
4
+ class 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 input
13
+ text = inputs.get("text", "")
14
+ inputs = self.tokenizer(text, return_tensors="pt").to(self.device)
15
+
16
+ # Inference
17
+ with torch.no_grad():
18
+ outputs = self.model.generate(**inputs)
19
+
20
+ # Post-process output
21
+ generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
22
+ return {"generated_text": generated_text}