nkasmanoff commited on
Commit
073c8a9
1 Parent(s): 0a479b8

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -5
app.py CHANGED
@@ -11,17 +11,17 @@ model = AutoModelForCausalLM.from_pretrained("nkasmanoff/git-planet").to(device)
11
 
12
  nlp = spacy.load('en_core_web_sm')
13
 
14
- def predict(image,max_length=64,device='cpu'):
15
  pixel_values = processor(images=image, return_tensors="pt").to(device).pixel_values
16
  generated_ids = model.generate(pixel_values=pixel_values, max_length=max_length)
17
  generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
18
- relation = get_relation(generated_caption)
19
 
20
- entity_pair = get_entities(generated_caption)
21
 
22
- knowlege_triplet = f"'{entity_pair[0]}'---{relation}--->'{entity_pair[1]}'"
23
 
24
- return knowlege_triplet
25
 
26
 
27
  def get_entities(sent):
 
11
 
12
  nlp = spacy.load('en_core_web_sm')
13
 
14
+ def predict(image,max_length=50,device='cpu'):
15
  pixel_values = processor(images=image, return_tensors="pt").to(device).pixel_values
16
  generated_ids = model.generate(pixel_values=pixel_values, max_length=max_length)
17
  generated_caption = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
18
+ # relation = get_relation(generated_caption)
19
 
20
+ # entity_pair = get_entities(generated_caption)
21
 
22
+ # knowlege_triplet = f"'{entity_pair[0]}'---{relation}--->'{entity_pair[1]}'"
23
 
24
+ return generated_caption
25
 
26
 
27
  def get_entities(sent):