Update handler.py

handler.py

@@ -1,14 +1,15 @@
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import torch
 
+MODEL_NAME = "teapotai/tinyteapot"
 MAX_INPUT_TOKENS = 512
 
+
 class EndpointHandler:
     def __init__(self, path: str = ""):
-        # IMPORTANT: on HF Endpoints, `path` is the mounted model directory (e.g. "/repository")
-        # Use it exactly, otherwise transformers can't find config.json/model files.
-        self.tokenizer = AutoTokenizer.from_pretrained(path)
-        self.model = AutoModelForSeq2SeqLM.from_pretrained(path)
+        # EXACTLY like your snippet: load by repo id from Hugging Face Hub
+        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)
 
         self.model.eval()
         self.device = torch.device("cpu")
@@ -29,7 +30,7 @@ class EndpointHandler:
             raise ValueError("Missing required field 'inputs'.")
 
         # Support:
-        # 1) inputs is a full prompt string (closest to your code)
+        # 1) inputs is a full prompt string (closest to your snippet)
         # 2) inputs is {"context": ..., "question": ...}
         if isinstance(inputs, str):
             prompt = inputs
@@ -40,12 +41,12 @@ class EndpointHandler:
         else:
             raise ValueError("inputs must be a string OR an object with 'context' and 'question'.")
 
-        # Match your code: tokenizer(prompt, return_tensors="pt")
+        # Match your code path
         enc = self.tokenizer(prompt, return_tensors="pt")
         input_ids = enc["input_ids"]
         attention_mask = enc["attention_mask"]
 
-        # Keep only most recent tokens (left-truncate)
+        # Keep only the most recent tokens if too long
         if input_ids.shape[1] > MAX_INPUT_TOKENS:
             input_ids = input_ids[:, -MAX_INPUT_TOKENS:]
             attention_mask = attention_mask[:, -MAX_INPUT_TOKENS:]
@@ -53,13 +54,11 @@ class EndpointHandler:
         input_ids = input_ids.to(self.device)
         attention_mask = attention_mask.to(self.device)
 
-        # Match your code: generate(..., do_sample=False)
         outputs = self.model.generate(
             input_ids=input_ids,
             attention_mask=attention_mask,
-            do_sample=False
+            do_sample=False,
         )
 
-        # Match your code: decode(..., skip_special_tokens=True)
         answer = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"generated_text": answer}