import requests
import json
import os

API_URL = "https://kp4xdy196cw81uf3.us-east-1.aws.endpoints.huggingface.cloud"
token = os.getenv('HUGGINGFACEHUB_API_TOKEN')
headers = {
	"Accept" : "application/json",
	"Authorization": "Bearer " + token,
	"Content-Type": "application/json"
}

def query(payload):
	response = requests.post(API_URL, headers=headers, json=payload)
	return response.json()

def generate_response(text):
  input = {
    "inputs": text,
    "parameters": {
      "max_new_tokens" : 256,
      "top_k": 10,
      "top_p": 0.95,
      "typical_p": 0.95,
      "temperature": 0.01,
      "repetition_penalty": 1.03,
      "stop" : ["/nHuman:", "/nUser:", "<end of message>\n", "Human:", "<end of response>"]
    }
  }
  
  output = query(input)
  return output[0]["generated_text"]