#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Sat Jun  1 17:30:57 2024
@author: air
"""

import streamlit as st
import requests
import os

# URL de l'API Hugging Face
API_URL = "https://api-inference.huggingface.co/models/condeac/condeac-Llama-2-7b-chat-finetune-medical"

# Récupérer le token d'accès Hugging Face depuis les variables d'environnement
hf_token = os.getenv("HF_API_TOKEN")
headers = {"Authorization": f"Bearer {hf_token}"}

# Limiter la longueur du prompt
MAX_PROMPT_LENGTH = 512

# Définir une fonction pour générer du texte en utilisant l'API
def generate_text(prompt, max_length=200):
    # Tronquer le prompt si nécessaire
    prompt = prompt[:MAX_PROMPT_LENGTH]
    
    payload = {
        "inputs": prompt,
        "parameters": {
            "max_length": max_length,
            "num_beams": 3,  # Réduire le nombre de faisceaux pour optimiser les ressources
            "no_repeat_ngram_size": 2,
            "early_stopping": True
        }
    }
    response = requests.post(API_URL, headers=headers, json=payload)
    response_json = response.json()
    
    # Afficher le contenu de la réponse pour le débogage
    print(response_json)
    
    # Gérer les erreurs et les structures inattendues de la réponse
    if response.status_code != 200:
        return f"Erreur: {response.status_code} - {response_json}"
    
    if isinstance(response_json, dict) and 'error' in response_json:
        return f"Erreur: {response_json['error']}"
    
    if isinstance(response_json, list) and len(response_json) > 0:
        return response_json[0].get('generated_text', 'Aucun texte généré')
    
    return "Réponse inattendue de l'API."

# Interface Streamlit
st.title("Chatbot Médical")
st.write("Posez une question ou entrez un prompt ci-dessous pour obtenir une réponse.")

# Zone de texte pour l'utilisateur
user_input = st.text_input("Vous : ", "hello !")

if st.button("Envoyer"):
    with st.spinner('en cours...'):
        response = generate_text(user_input)
        st.success("Chatbot : " + response)