Update app.py

app.py

@@ -19,7 +19,11 @@ base_model_id = "codellama/CodeLlama-7b-hf"
 dataset = "gretelai/synthetic_text_to_sql"
 
 def prompt_model(model_id, system_prompt, user_prompt, schema):
-    pipe = pipeline("text-generation", model=model_id, model_kwargs={"torch_dtype": torch.bfloat16}, device_map="auto")
+    pipe = pipeline("text-generation", 
+                    model=model_id, 
+                    model_kwargs={"torch_dtype": torch.bfloat16}, 
+                    device_map="auto",
+                    max_new_tokens=1000)
     messages = [
       {"role": "system", "content": system_prompt.format(schema=schema)},
       {"role": "user", "content": user_prompt},