import os
# os.system('pip install transformers')

import gradio as gr
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin


def greet(name):
    return "Hello " + name + "!!"

class MyModel(nn.Module, PyTorchModelHubMixin):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(20,50)

    def forward(self, x):
        return x+2
        
model = MyModel.from_pretrained('WeixuanYuan/test_upload_private', use_auth_token='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')

def plus_two(input):
    return model(int(input))

iface = gr.Interface(fn=plus_two, inputs="text", outputs="text")
iface.launch()