Update app.py

app.py

@@ -1,13 +1,13 @@
 import gradio as gr
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-model_name = "bigcode/starcoderbase"
+model_name = "Salesforce/codegen-350M-multi"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
 def generate_code(prompt):
     inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
+    outputs = model.generate(**inputs, max_new_tokens=128)
     code = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return code
 
@@ -19,4 +19,3 @@ iface = gr.Interface(
 )
 
 iface.launch()
-