Git (Custom Code)

import torch
from transformers import pipeline
from pydantic import BaseModel, Field
import inferless

@inferless.request
class RequestObjects(BaseModel):
    prompt: str = Field(default="a horse near a beach")

@inferless.response
class ResponseObjects(BaseModel):
    generated_txt: str = Field(default='Test output')

app = inferless.Cls(gpu="T4")
class InferlessPythonModel:
    @app.load
    def initialize(self):
        self.generator = pipeline("text-generation", model="EleutherAI/gpt-neo-125M",device=0)
    @app.infer
    def infer(self, inputs):
        pipeline_output = self.generator(inputs.prompt, do_sample=True, min_length=128)
        generateObject = ResponseObjects(generated_txt = pipeline_output[0]["generated_text"])
        return generate object

## Implement the Load function here for the model 
    def initialize(self):
        self.generator = pipeline("text-generation", model="EleutherAI/gpt-neo-125M",device=0)

    
# Function to perform inference 
    def infer(self, inputs):
        # inputs is a dictionary where the keys are input names and values are actual input data
        # e.g. in the below code the input name is "prompt"
        prompt = inputs["prompt"]
        pipeline_output = self.generator(prompt, do_sample=True, min_length=20)
        generated_txt = pipeline_output[0]["generated_text"]
        # The output generated by the infer function should be a dictionary where keys are output names and values are actual output data
        # e.g. in the below code the output name is "generated_txt"
        return {"generated_text": generated_txt}

# perform any cleanup activity here
    def finalize(self,args):
        self.pipe = None

INPUT_SCHEMA = {
    "prompt": {
        'datatype': 'STRING',
        'required': True,
        'shape': [1],
        'example': ["There is a fine house in the forest"]
    }
}