from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import os
import requests

HF_TOKEN = os.environ.get('HF_TOKEN')

app = FastAPI(title="Deploying FastAPI Apps on Huggingface")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


API_URL = "https://api-inference.huggingface.co/models/google/gemma-1.1-7b-it"
headers = {"Authorization": f"Bearer {HF_TOKEN}"}

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

@app.get("/query", tags=["Query"])
def api_query(input: str):
    return query({"inputs": input,"parameters":{"return_full_text":False,"max_length":1024}})

@app.get("/", tags=["Home"])
def api_home():
    return {'message': 'Welcome!'}