Upload inference.py

inference.py

@@ -0,0 +1,97 @@
+import os
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+# Specify the model directory (SageMaker uses /opt/ml/model by default)
+def model_fn(model_dir):
+    import os
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+def model_fn(model_dir):
+    """
+    Load the model and tokenizer from the specified directory.
+    """
+    print("Loading model and tokenizer...")
+
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_dir)
+
+    # Explicitly load the .pth file
+    model_path = os.path.join(model_dir, "pytorch_model.pth")
+    config_path = os.path.join(model_dir, "config.json")
+
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model file not found: {model_path}")
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    # Load the model using the state dictionary
+    config = AutoConfig.from_pretrained(config_path)
+    model = AutoModelForSeq2SeqLM(config)
+    model.load_state_dict(torch.load(model_path))
+    print("Model and tokenizer loaded successfully.")
+
+    return model, tokenizer
+
+def input_fn(serialized_input_data, content_type="application/json"):
+    """
+    Deserialize the input data from JSON format.
+    """
+    print("Processing input data...")
+    if content_type == "application/json":
+        import json
+        input_data = json.loads(serialized_input_data)
+        if "plsql_code" not in input_data:
+            raise ValueError("Missing 'plsql_code' in the input JSON.")
+        print("Input data processed successfully.")
+        return input_data["plsql_code"]
+    else:
+        raise ValueError(f"Unsupported content type: {content_type}")
+
+def predict_fn(input_data, model_and_tokenizer):
+    """
+    Translate PL/SQL code to Hibernate/JPA-based Java code using the trained model.
+    """
+    print("Starting prediction...")
+    model, tokenizer = model_and_tokenizer
+
+    # Construct the tailored prompt
+    prompt = f"""
+    Translate this PL/SQL function to a Hibernate/JPA-based Java implementation.
+
+    Requirements:
+    1. Use @Entity, @Table, and @Column annotations to map the database table structure.
+    2. Define Java fields corresponding to the database columns used in the PL/SQL logic.
+    3. Replicate the PL/SQL logic as a @Query in the Repository layer or as Java logic in the Service layer.
+    4. Use Repository and Service layers, ensuring transactional consistency with @Transactional annotations.
+    5. Avoid direct bitwise operations in procedural code; ensure they are part of database entities or queries.
+
+    Input PL/SQL:
+    {input_data}
+    """
+    # Tokenize and generate the translated code
+    print("Tokenizing input...")
+    inputs = tokenizer(prompt, return_tensors="pt", max_length=1024, truncation=True)
+    print("Generating output...")
+    outputs = model.generate(
+        inputs["input_ids"],
+        max_length=1024,
+        num_beams=4,
+        early_stopping=True
+    )
+    translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    print("Prediction completed.")
+    return translated_code
+
+def output_fn(prediction, accept="application/json"):
+    """
+    Serialize the prediction result to JSON format.
+    """
+    print("Serializing output...")
+    if accept == "application/json":
+        import json
+        return json.dumps({"translated_code": prediction}), "application/json"
+    else:
+        raise ValueError(f"Unsupported accept type: {accept}")
+