Update app.py

app.py

@@ -3,12 +3,73 @@ import torch
 import json
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
+if torch.cuda.is_available():
+    use_cuda = True
+else:
+    use_cuda = False
+
 tokenizer = AutoTokenizer.from_pretrained("keminglu/pivoine-7b", use_auth_token="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", padding_side="left")
 model = AutoModelForCausalLM.from_pretrained("keminglu/pivoine-7b", use_auth_token="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", torch_dtype=torch.float16)
-#input_device = torch.device("cuda:5")
 model.requires_grad_(False)
 model.eval()
-#model = model.to(input_device)
+if use_cuda:
+    model = model.to("cuda")
 
 examples = json.load(open("examples.json"))
-description = open("description.txt").read()
+description = open("description.txt").read()
+
+def inference(context, instruction, num_beams:int=4):
+    input_str = f"\"{context}\"\n\n{instruction}"
+    if not input_str.endswith("."):
+        input_str += "."
+
+    input_tokens = tokenizer(input_str, return_tensors="pt", padding=True)
+    if use_cuda:
+        for t in input_tokens:
+            if torch.is_tensor(input_tokens[t]):
+                input_tokens[t] = input_tokens[t].to(input_device)
+
+    output = model.generate(
+        input_tokens['input_ids'],
+        num_beams=num_beams,
+        do_sample=False,
+        max_new_tokens=2048,
+        num_return_sequences=1,
+        return_dict_in_generate=True,
+    )
+
+    num_input_tokens = input_tokens["input_ids"].shape[1]
+    output_tokens = output.sequences
+    generated_tokens = output_tokens[:, num_input_tokens:]
+    num_generated_tokens = (generated_tokens != tokenizer.pad_token_id).sum(dim=-1).tolist()[0]
+    prefix_to_add = torch.tensor([[tokenizer("A")["input_ids"][0]]]).to(input_device)
+    generated_tokens = torch.cat([prefix_to_add, generated_tokens], dim=1)
+    generated_text = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    string_output = [i[1:].strip() for i in generated_text][0]
+    json_output = None
+    try:
+        json_output = json.loads(string_output)
+    except json.JSONDecodeError:
+        json_output = {"error": "Unfortunately, there is a JSON decode error on your output, which is really rare in our experiment :("}
+    except Exception as e:
+        raise gr.Error(e)
+
+    return num_generated_tokens, string_output, json_output
+
+demo = gr.Interface(
+        fn=inference,
+        inputs=["text", "text", gr.Slider(1,5,value=4,step=1)],
+        outputs=[
+            gr.Number(label="Number of Generated Tokens"),
+            gr.Textbox(label="Raw String Output"),
+            gr.JSON(label="Json Output")],
+        examples=examples,
+        examples_per_page=3,
+        title="Instruction-following Open-world Information Extraction",
+        description=description,
+    )
+
+demo.launch(
+    share=True,
+    auth=("miega", "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"),
+    show_error=True)