import gradio as gr
from transformers import pipeline
import torch

# Check if a GPU is available
device = 0 if torch.cuda.is_available() else -1

# Load the text-generation pipeline with the appropriate device
model = pipeline(
    "text-generation",
    model="rish13/polymers",
    device=device  # Automatically use GPU if available, otherwise CPU
)

def generate_response(prompt):
    # Generate text from the model
    response = model(
        prompt, 
        max_length=50,  # Adjusted to generate shorter text
        num_return_sequences=1, 
        temperature=0.5,  # Lowered to reduce randomness
        top_k=50,  # Limiting the next word selection
        top_p=0.9  # Cumulative probability threshold
    )
    
    # Get the generated text from the response
    generated_text = response[0]['generated_text']
    
    return generated_text

# Define the Gradio interface
interface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(
        lines=2, 
        placeholder="Enter your prompt here...", 
        label="Prompt",
        elem_id="input-textbox"  # Custom styling ID for input textbox
    ),
    outputs=gr.Textbox(
        label="Generated Text",
        elem_id="output-textbox"  # Custom styling ID for output textbox
    ),
    title="Polymer Knowledge Model",
    description=(
        "This application uses a fine-tuned model to generate text related to polymers. "
        "Enter a prompt to get started, and the model will generate relevant text."
    ),
    theme="huggingface",  # Apply a theme for consistent styling
    layout="horizontal",  # Arrange input and output side by side
    live=True  # Update the output live as the user types
)

# Launch the interface
interface.launch()