Update app.py

app.py

@@ -6,7 +6,11 @@ import torch
 # Load the model and config when the script starts
 config = PeftConfig.from_pretrained("phearion/bigbrain-v0.0.1")
 model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-7B-v0.1")
-model = PeftModel.from_pretrained(model, "phearion/bigbrain-v0.0.1")
+model = PeftModel.from_pretrained(model,
+ "phearion/bigbrain-v0.0.1")
+
+# Convert the model to TorchScript
+scripted_model = torch.jit.script(model)
 
 # Load the tokenizer
 tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
@@ -18,7 +22,7 @@ def greet(text):
     
     # Use torch.no_grad to disable gradient calculation
     with torch.no_grad():
-        output_tokens = model.generate(**batch, max_new_tokens=20)
+        output_tokens = scripted_model.generate(**batch, max_new_tokens=20)
 
     return tokenizer.decode(output_tokens[0], skip_special_tokens=True)