app.py

import gradio as gr
from transformers import AutoModelWithLMHead, AutoTokenizer
import torch

# Load your model and tokenizer from the hub
model = AutoModelWithLMHead.from_pretrained("shrivasc/my-finetuned-gpt2-medquad")  # Replace with your model path
tokenizer = AutoTokenizer.from_pretrained("shrivasc/my-finetuned-gpt2-medquad")  # Replace with your tokenizer path


def generate_response(prompt, max_length=200):
    """
    Generates a response using the fine-tuned GPT-2 model.

    Args:
        prompt (str): The input prompt text.
        max_length (int, optional): The maximum length of the generated response. Defaults to 200.

    Returns:
        str: The generated response text.
    """
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    if max_length is None:
        max_length = len(input_ids[0]) + 1

    # Check the device of the model
    device = next(model.parameters()).device

    # Move input_ids to the same device as the model
    input_ids = input_ids.to(device)

    # Create the attention mask and pad token id
    attention_mask = torch.ones_like(input_ids)
    pad_token_id = tokenizer.eos_token_id

    output = model.generate(
        input_ids,
        max_length=100,
        num_return_sequences=1,
        attention_mask=attention_mask,
        pad_token_id=pad_token_id,
        temperature=0.7,  # Adjust the temperature value
        repetition_penalty=1.2
    )

    return tokenizer.decode(output[0], skip_special_tokens=True)


# Create input and output Gradio elements
in_prompt = gr.Textbox(label="Enter your query")
out_response = gr.Textbox(label="Response")

# Create a Gradio interface object
iface = gr.Interface(fn=generate_response, inputs=[in_prompt], outputs=out_response)

# Launch the interface to generate UI
iface.launch()