import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer

# Load your fine-tuned GPT-2 model from Hugging Face
MODEL_NAME = "hackergeek98/finetuned-gpt2"  # Replace with your model
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# Function to generate responses
def generate_response(user_input):
    # Tokenize the input
    inputs = tokenizer(user_input, return_tensors="pt")

    # Generate a response
    outputs = model.generate(inputs['input_ids'], max_length=1000, num_return_sequences=1, no_repeat_ngram_size=2)

    # Decode the output and return the result
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response

# Create Gradio interface
interface = gr.Interface(fn=generate_response,
                         inputs=gr.Textbox(label="Enter your message"),
                         outputs=gr.Textbox(label="Therapist Response"),
                         title="Virtual Therapist",
                         description="A fine-tuned GPT-2 model acting as a virtual therapist.")

# Launch the app
interface.launch()