| import gradio as gr |
| from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
| def predict_code(input): |
| tokenizer = AutoTokenizer.from_pretrained('JetBrains/Mellum-4b-base') |
| model = AutoModelForCausalLM.from_pretrained('JetBrains/Mellum-4b-base') |
| encoded_input = tokenizer(input, return_tensors='pt', return_token_type_ids=False) |
| input_len = len(encoded_input["input_ids"][0]) |
| out = model.generate( |
| **encoded_input, |
| max_new_tokens=100, |
| ) |
|
|
| prediction = tokenizer.decode(out[0][input_len:]) |
| return prediction |
|
|
| def run(input): |
| return predict_code(input) |
|
|
| app = gr.Interface( |
| fn=run, |
| inputs=["text"], |
| outputs=["text"] |
| ) |
|
|
| app.launch() |
|
|
