import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

# Load the SmolLM-135M model and tokenizer
model_name = "HuggingFaceTB/SmolLM-135M-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Initialize feedback history
feedback_history = []

def generate_unique_feedback(text):
    # Generate a prompt for writing feedback
    prompt = f"Provide constructive feedback on the following creative writing piece:\n\n{text}\n\nFeedback:"
    
    # Tokenize the input with attention mask
    inputs = tokenizer(prompt, return_tensors="pt", padding=True)
    
    # Attempt to generate unique feedback
    for _ in range(5):  # Try up to 5 times to get unique feedback
        with torch.no_grad():
            # Use max_new_tokens instead of max_length
            outputs = model.generate(inputs.input_ids, attention_mask=inputs.attention_mask, max_new_tokens=300, do_sample=True, top_p=0.85, temperature=0.7)
        
        # Decode the response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        feedback = response.split("Feedback:")[-1].strip()
        
        # Check if feedback is unique
        if feedback not in feedback_history:
            feedback_history.append(feedback)  # Add new feedback to history
            return feedback

    # If all generated feedback is repeated, return a default message
    return "No new feedback available at this time. Try rephrasing or adding more details to your text."

# Reset history function for Gradio button
def reset_history():
    global feedback_history
    feedback_history = []  # Clear the history
    return "Feedback history has been reset."

# Set up Gradio interface using Blocks
with gr.Blocks() as app:
    gr.Markdown("# WriteBetter\nQuick feedback on tone, grammar, and word choice.")
    
    with gr.Row():
        input_text = gr.Textbox(lines=5, label="Your Writing", placeholder="Paste a short piece of creative writing here...")
        feedback_output = gr.Textbox(label="Feedback")
    
    feedback_button = gr.Button("Generate Feedback")
    feedback_button.click(fn=generate_unique_feedback, inputs=input_text, outputs=feedback_output)
    
    with gr.Row():
        reset_status = gr.Textbox(label="Reset Status", interactive=False)
    
    reset_button = gr.Button("Reset Feedback History")
    reset_button.click(fn=reset_history, outputs=reset_status)

# Launch the app with a public shareable link
app.launch(share=True)