DisgustingOzil commited on
Commit
17970fe
·
verified ·
1 Parent(s): 257e5a2

create handler.[y

Browse files
Files changed (1) hide show
  1. handler.py +22 -0
handler.py ADDED
@@ -0,0 +1,22 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Dict, List, Any
2
+ from transformers import AutoTokenizer, AutoModelForCausalLM
3
+ import holidays
4
+
5
+ class EndpointHandler():
6
+ def __init__(self, path=""):
7
+ model_id = "DisgustingOzil/Academic-MCQ-Generator"
8
+ load_in_4bit = True
9
+ max_seq_length = 2048
10
+ tokenizer = AutoTokenizer.from_pretrained(model_id)
11
+ model = AutoModelForCausalLM.from_pretrained(model_id, load_in_4bit=load_in_4bit)
12
+
13
+ def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
14
+ input_text = data.pop("input_text",data)
15
+ inputs = tokenizer(input_text, return_tensors="pt")
16
+ outputs = model.generate(
17
+ **inputs,
18
+ max_length=1000,
19
+ num_return_sequences=1,
20
+ )
21
+ output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
22
+ return output_text