adding story func

app.py

@@ -13,6 +13,21 @@ import torch
 # https://github.com/NielsRogge/Transformers-Tutorials/blob/master/HuggingFace_vision_ecosystem_overview_(June_2022).ipynb
 # option 1: load with randomly initialized weights (train from scratch)
 
+from transformers import (
+    AutoModelForCausalLM,
+    LogitsProcessorList,
+    MinLengthLogitsProcessor,
+    StoppingCriteriaList,
+    MaxLengthCriteria,
+)
+
+
+
+
+
+
+
+
 config = ViTConfig(num_hidden_layers=12, hidden_size=768)
 model = ViTForImageClassification(config)
 
@@ -28,6 +43,36 @@ model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/
 #google/vit-base-patch16-224, deepmind/vision-perceiver-conv
 image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor)
 
+def create_story(text_seed):
+  tokenizer = AutoTokenizer.from_pretrained("gpt2")
+  model = AutoModelForCausalLM.from_pretrained("gpt2")
+
+  # set pad_token_id to eos_token_id because GPT2 does not have a EOS token
+  model.config.pad_token_id = model.config.eos_token_id
+
+  #input_prompt = "It might be possible to"
+  input_prompt = text_seed
+  input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids
+
+  # instantiate logits processors
+  logits_processor = LogitsProcessorList(
+    [
+        MinLengthLogitsProcessor(10, eos_token_id=model.config.eos_token_id),
+    ]
+  )
+  stopping_criteria = StoppingCriteriaList([MaxLengthCriteria(max_length=20)])
+
+  outputs = model.greedy_search(
+    input_ids, logits_processor=logits_processor, stopping_criteria=stopping_criteria
+  )
+
+  result_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)
+  return result_text
+
+
+
+
+
 
 def self_caption(image):
   repo_name = "ydshieh/vit-gpt2-coco-en"
@@ -60,7 +105,8 @@ def self_caption(image):
   print(pred_dictionary)
   #return(pred_dictionary)
   preds = ' '.join(preds)
-  return preds
+  story = create_story(preds)
+  return story
 
 
 def classify_image(image):