app

app.py

@@ -0,0 +1,40 @@
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+hf_model = "law-llm/law-glm-10b"
+max_question_length = 64
+max_generation_length = 490
+
+tokenizer = AutoTokenizer.from_pretrained(
+    hf_model,
+    cache_dir=model_cache_dir,
+    use_fast=True,
+    trust_remote_code=True
+)
+
+model = AutoModelForSeq2SeqLM.from_pretrained(
+    hf_model,
+    cache_dir=model_cache_dir,
+    trust_remote_code=True
+)
+
+model = model.to('cuda')
+model.eval()
+
+model_inputs = "提问: 犯了盗窃罪怎么判刑? 回答: [gMASK]"
+
+model_inputs = tokenizer(model_inputs,
+                            max_length=max_question_length,
+                            padding=True,
+                            truncation=True,
+                            return_tensors="pt")
+
+model_inputs = tokenizer.build_inputs_for_generation(model_inputs,
+                                                     targets=None,
+                                                     max_gen_length=max_generation_length,
+                                                     padding=True)
+
+inputs = model_inputs.to('cuda')
+
+outputs = model.generate(**inputs, max_length=max_generation_length,
+                            eos_token_id=tokenizer.eop_token_id)
+prediction = tokenizer.decode(outputs[0].tolist())