app.py

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

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("braindeck/text2text", trust_remote_code=True, subfolder="checkpoints/model")
model = AutoModelForCausalLM.from_pretrained("braindeck/text2text", trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto", subfolder="checkpoints/model")

def generate_response(prompt):
    """
    Generates a response from the model.
    """
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=512)
    
    # Decode the generated text
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
    return generated_text

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Text-to-Text Generation with DeepSeek-R1-Distill-Qwen-7B")
    gr.Markdown("Enter a prompt and the model will generate a response.")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Prompt", lines=4, placeholder="Enter your prompt here...")
        
    with gr.Row():
        generate_button = gr.Button("Generate")
        
    with gr.Row():
        response_output = gr.Textbox(label="Response", lines=8, interactive=False)
        
    generate_button.click(
        fn=generate_response,
        inputs=prompt_input,
        outputs=response_output
    )

if __name__ == "__main__":
    demo.launch()