import streamlit as st
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

# Load the fine-tuned GPT-2 model and tokenizer
model_name = "marianeft/MedQuAD"
model = GPT2LMHeadModel.from_pretrained(model_name)
tokenizer = GPT2Tokenizer.from_pretrained(model_name)

# Ensure tokenizer has a pad token
tokenizer.pad_token = tokenizer.eos_token

st.title("Medical Query Response Generator")
st.write("This app generates responses to medical queries using a fine-tuned GPT-2 model.")

# User input for medical query
query = st.text_area("Enter your medical query:")

# User input for response length
response_length = st.slider("Select response length:", min_value=10, max_value=300, value=50)

# Generate response when user submits query
if st.button("Generate Response"):
    if query.strip():
        # Encode input query and generate response
        inputs = tokenizer.encode(query, return_tensors="pt")
        outputs = model.generate(
            inputs,
            max_length=response_length,
            num_return_sequences=1,
            pad_token_id=tokenizer.eos_token_id  # Prevent padding errors
        )

        # Decode and display the response
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        st.subheader("Response:")
        st.write(response)
    else:
        st.warning("Please enter a query to generate a response.")