import os
from dotenv import load_dotenv
import requests
import json
import gradio as gr

# Load environment variables
load_dotenv()

# Access the Hugging Face API key and endpoint URL
hf_api_key = os.getenv('HF_API_KEY')
MODEL_NAME = "sshleifer/distilbart-cnn-12-6"
ENDPOINT_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"

def get_completion(inputs, parameters=None):
    headers = {
        "Authorization": f"Bearer {hf_api_key}",
        "Content-Type": "application/json"
    }
    data = {
        "inputs": inputs
    }
    if parameters is not None:
        data.update({"parameters": parameters})
    
    try:
        response = requests.post(ENDPOINT_URL, headers=headers, json=data)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Request failed: {e}")
        return {"error": f"Request failed: {str(e)}"}

def summarize(input_text):
    try:
        output = get_completion(input_text)
        if isinstance(output, list) and len(output) > 0 and 'summary_text' in output[0]:
            return output[0]['summary_text']
        else:
            return f"Unexpected response format: {output}"
    except Exception as e:
        return f"An error occurred: {str(e)}"

demo = gr.Interface(fn=summarize, inputs="text", outputs="text")
demo.launch()