import gradio as gr
from transformers import pipeline

# Load the model with optimization settings
model = pipeline(
    "text-generation",
    model="rish13/polymers",
    device=0  # Use device=0 for GPU (if available), -1 for CPU
)

def generate_response(prompt):
    # Generate text from the model
    response = model(prompt, max_length=100, num_return_sequences=1, temperature=0.7)
    
    # Get the generated text from the response
    generated_text = response[0]['generated_text']

    # Find the position of the first end-of-sentence punctuation
    end_punctuation = ['.', '!', '?']
    end_position = -1
    for punct in end_punctuation:
        pos = generated_text.find(punct)
        if pos != -1 and (end_position == -1 or pos < end_position):
            end_position = pos

    # If punctuation is found, truncate the text at that point
    if end_position != -1:
        generated_text = generated_text[:end_position + 1]
    
    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"),
    outputs="text",
    title="Polymer Knowledge Model",
    description="A model fine-tuned for generating text related to polymers."
)

# Launch the interface
interface.launch()