import streamlit as st
import os

from huggingface_hub import InferenceClient

# Get token from Space secrets
API_TOKEN = os.getenv("HF_TOKEN")

# Initialize the Inference Client
client = InferenceClient(token=API_TOKEN)

# Title of the app
st.title("Sentence Improver App")

# Text input from the user
user_input = st.text_input("Enter a sentence to improve:", "I goed to the park and play.")

# Button to trigger the correction
if st.button("Improve Sentence"):
    if user_input:
        # Create a prompt for the LLM to improve the sentence
        prompt = f"Correct and improve this sentence: '{user_input}'"
        
        # Call the LLM via Hugging Face Inference API
        try:
            response = client.text_generation(
                prompt,
                model="mistralai/Mixtral-8x7B-Instruct-v0.1",  # You can change this to another model
                #model="deepseek-ai/deepseek-coder-6.7b-instruct",
                max_new_tokens=100,
                temperature=0.7,
            )
            # Extract the improved sentence (remove the prompt part if included)
            improved_sentence = response.strip()
            st.write("Improved Sentence:", improved_sentence)
        except Exception as e:
            st.error(f"Error calling the LLM: {str(e)}")
    else:
        st.warning("Please enter a sentence first!")
else:
    st.write("Enter a sentence and click the button to see it improved!")