Update app.py

app.py

@@ -16,6 +16,26 @@ from transformers import (
     StoppingCriteriaList,
     MaxLengthCriteria,
 )
+import json
+import os
+from screenshot import (
+    before_prompt,
+    prompt_to_generation,
+    after_generation,
+    js_save,
+    js_load_script,
+)
+from spaces_info import description, examples, initial_prompt_value
+
+API_URL = os.getenv("API_URL")
+HF_API_TOKEN = os.getenv("HF_API_TOKEN")
+
+
+def query(payload):
+    print(payload)
+    response = requests.request("POST", API_URL, json=payload, headers={"Authorization": f"Bearer {HF_API_TOKEN}"})
+    print(response)
+    return json.loads(response.content.decode("utf-8"))
 
 # 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)
@@ -39,6 +59,52 @@ model = PerceiverForImageClassificationConvProcessing.from_pretrained("deepmind/
 #google/vit-base-patch16-224, deepmind/vision-perceiver-conv
 image_pipe = ImageClassificationPipeline(model=model, feature_extractor=feature_extractor)
 
+
+
+
+def inference(input_sentence, max_length, sample_or_greedy, seed=42):
+    if sample_or_greedy == "Sample":
+        parameters = {
+            "max_new_tokens": max_length,
+            "top_p": 0.9,
+            "do_sample": True,
+            "seed": seed,
+            "early_stopping": False,
+            "length_penalty": 0.0,
+            "eos_token_id": None,
+        }
+    else:
+        parameters = {
+            "max_new_tokens": max_length,
+            "do_sample": False,
+            "seed": seed,
+            "early_stopping": False,
+            "length_penalty": 0.0,
+            "eos_token_id": None,
+        }
+
+    payload = {"inputs": input_sentence, "parameters": parameters,"options" : {"use_cache": False} }
+
+    data = query(payload)
+
+    if "error" in data:
+        return (None, None, f"<span style='color:red'>ERROR: {data['error']} </span>")
+
+    generation = data[0]["generated_text"].split(input_sentence, 1)[1]
+    return (
+        before_prompt
+        + input_sentence
+        + prompt_to_generation
+        + generation
+        + after_generation,
+        data[0]["generated_text"],
+        "",
+    )
+
+
+
+
+
 def create_story(text_seed):
   #tokenizer = AutoTokenizer.from_pretrained("gpt2")
   #model = AutoModelForCausalLM.from_pretrained("gpt2")