Create inference.py

inference.py

@@ -1,63 +1,13 @@
-import torch
-from transformers import T5Tokenizer
-from model import GPT
-
-class Inference:
-    def __init__(self, model_path, tokenizer_path, device='cuda' if torch.cuda.is_available() else 'cpu'):
-        self.device = device
-        self.tokenizer = T5Tokenizer.from_pretrained(tokenizer_path)
-        self.model = GPT(
-            vocab_size=self.tokenizer.vocab_size,
-            embed_size=1500,
-            num_layers=20,
-            heads=20,
-            expansion_factor=4,
-            dropout=0.1,
-            max_length=1024
-        )
-        self.model.load_state_dict(torch.load(model_path, map_location=self.device))
-        self.model.to(self.device)
-        self.model.eval()
-
-    def predict(self, text, max_length=100):
-        input_ids = self.tokenizer.encode(text, return_tensors='pt').to(self.device)
-        generated_tokens = set(input_ids[0].tolist())
-
-        with torch.no_grad():
-            for _ in range(max_length):
-                outputs = self.model(input_ids)
-                logits = outputs[:, -1, :] / 1.0  # temperature = 1.0
-
-                for token_id in generated_tokens:
-                    logits[0, token_id] /= 1.5  # repetition_penalty = 1.5
-
-                filtered_logits = top_k_top_p_filtering(logits, top_k=50, top_p=0.9)
-                probs = torch.softmax(filtered_logits, dim=-1)
-
-                next_token_id = torch.multinomial(probs, 1)
-                next_token_id = next_token_id.squeeze(-1).unsqueeze(0)
-                input_ids = torch.cat([input_ids, next_token_id], dim=1)
-
-                generated_tokens.add(next_token_id.item())
-
-                if next_token_id.item() == self.tokenizer.eos_token_id:
-                    break
-
-        return self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
-
-def top_k_top_p_filtering(logits, top_k=0, top_p=0.9, filter_value=-float('Inf')):
-    top_k = min(top_k, logits.size(-1))
-    if top_k > 0:
-        indices_to_remove = logits < torch.topk(logits, top_k).values[:, -1, None]
-        logits[indices_to_remove] = filter_value
-
-    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
-    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
-    sorted_indices_to_remove = cumulative_probs > top_p
-
-    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
-    sorted_indices_to_remove[..., 0] = 0
-
-    indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
-    logits[indices_to_remove] = filter_value
-    return logits
+from inference import Inference
+import os
+
+model_path = os.getenv("MODEL_PATH", "saved_model/pytorch_model.bin")
+tokenizer_path = os.getenv("TOKENIZER_PATH", "saved_tokenizer")
+inference = Inference(model_path, tokenizer_path)
+
+def handler(event, context):
+    prompt = event["data"]["prompt"]
+    max_length = event["data"].get("max_length", 100)
+    
+    response = inference.predict(prompt, max_length)
+    return {"response": response}