import os
import gradio as gr
from transformers import pipeline
from huggingface_hub import login

# Read the token from the environment variable
HUGGINGFACE_TOKEN = os.getenv("HUGGINGFACE_TOKEN")

# Authenticate with Hugging Face
if HUGGINGFACE_TOKEN:
    login(token=HUGGINGFACE_TOKEN)
else:
    raise ValueError("Hugging Face token not found in environment variables.")

# Initialize the text generation pipeline
pipe = pipeline("text-generation", model="TinyLlama/TinyLlama-1.1B-Chat-v1.0")

def generate_response(user_input):
    # Generate text based on the user's input
    response = pipe(user_input, max_length=100, num_return_sequences=1)
    # Extract the generated text
    generated_text = response[0]['generated_text']
    return generated_text

# Create the Gradio interface
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(label="prompt:", lines=2, placeholder="prompt"),
    outputs="text",
    title="Gemma",
    description="Prompt gemma-2b"
)

# Launch the Gradio app
interface.launch()