Update app.py

app.py

@@ -39,17 +39,17 @@ def conditional_compile(func):
         compiled = torch.compile(func, backend="openxla")
         return compiled
     return func
-        
 
+model_id = "google/gemma-2b"
+torch_dtype = torch.bfloat16
+model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype)
+device = model.device
+model = model.eval()
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
-def summarize(inp):
-    model_id = "google/gemma-2b"
-    torch_dtype = torch.bfloat16
-    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype)
-    device = model.device
-    model = model.eval()
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+def summarize(inp, model=model, tokenizer=tokenizer, device=device):
     with torch.no_grad():
         inp = inp.replace('\n','')
         inputs = tokenizer(inp, return_tensors="pt", padding=True).to(device)
@@ -106,3 +106,4 @@ def summarize(inp):
         return response
 
 gr.Interface(fn=summarize, inputs=gr.Textbox(lines=7, label="Input Text"), outputs="text", title="gemma-2b Demo").launch(inline=False)
+