Update app.py

app.py

@@ -1,8 +1,21 @@
 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 "Rofl: " + name
+    return predict_code(input)
 
 app = gr.Interface(
     fn=run,