import os
from flask import Flask, request, render_template
import requests

app = Flask(__name__)

API_URL = "https://api-inference.huggingface.co/models/TinyLlama/TinyLlama-1.1B-Chat-v1.0"
api_key = os.environ.get('HUGGING_FACE_API_KEY')  # Retrieve API key from environment variable

def query(payload):
  headers = {"Authorization": f"Bearer {api_key}"}
  response = requests.post(API_URL, headers=headers, json=payload)
  return response.json()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/get_response', methods=['POST'])
def get_response():
    user_input = request.form['user_input']
    output = query({"inputs": user_input})
    return output

if __name__ == '__main__':
    app.run(debug=True)