Update app.py

app.py

@@ -9,23 +9,13 @@ from pymatgen.io.ase import AseAtomsAdaptor
 from ase.io import write as ase_write
 import tempfile
 import time
-# 设置PyTorch使用的线程数
-torch.set_num_threads(2)
-def load_quantized_model(model_path):
-    model = MatterGPTWrapper.from_pretrained(model_path)
-    model.to('cpu')
-    model.eval()
-    quantized_model = torch.quantization.quantize_dynamic(
-        model, {torch.nn.Linear}, dtype=torch.qint8
-    )
-    return quantized_model
 
 
-# Load and quantize the model
+# Load the model
 model_path = "./"
-quantized_model = load_quantized_model(model_path)
-quantized_model.to("cpu")
-quantized_model.eval()
+model = MatterGPTWrapper.from_pretrained(model_path)
+model.to("cpu")
+model.eval()
 # Load the tokenizer
 tokenizer_path = "Voc_prior"
 tokenizer = SimpleTokenizer(tokenizer_path)
@@ -37,23 +27,22 @@ try:
 except Exception as e:
     backend = SLICES(relax_model=None)
 
-
-
-def generate_slices_quantized(quantized_model, tokenizer, formation_energy, band_gap, max_length, temperature, do_sample, top_k, top_p):
+def generate_slices_optimized(model, tokenizer, formation_energy, band_gap, max_length, temperature, do_sample, top_k, top_p):
     condition = torch.tensor([[float(formation_energy), float(band_gap)]], dtype=torch.float32)
     context = '>'
     x = torch.tensor([[tokenizer.stoi[context]]], dtype=torch.long)
     
     with torch.no_grad():
-        generated = quantized_model.generate(x, prop=condition, max_length=max_length, 
-                                             temperature=temperature, do_sample=do_sample, 
-                                             top_k=top_k, top_p=top_p)
+        generated = model.generate(x, prop=condition, max_length=max_length, 
+                                   temperature=temperature, do_sample=do_sample, 
+                                   top_k=top_k, top_p=top_p)
     
     return tokenizer.decode(generated[0].tolist())
 
 def generate_slices(formation_energy, band_gap):
-    return generate_slices_quantized(quantized_model, tokenizer, formation_energy, band_gap, 
-                                     quantized_model.config.block_size, 1.2, True, 0, 0.9)
+    return generate_slices_optimized(model, tokenizer, formation_energy, band_gap, 
+                                     model.config.block_size, 1.2, True, 0, 0.9)
+
 def wrap_structure(structure):
     """Wrap all atoms back into the unit cell."""
     for i, site in enumerate(structure):