Pardeep0675's picture
Upload 2 files
b0af655 verified
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch
# 1. SETUP MODEL
# We use a specialized T5 model for high-quality rewriting
model_name = "Vamsi/T5_Paraphrase_Paws"
device = "cpu" # Force CPU for free tier stability
print("Loading AI Model... this happens only once.")
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
def paraphrase_text(text):
if not text or len(text.strip()) < 5:
return "Please enter a longer sentence to rewrite."
try:
# Prepare text for the model
text = "paraphrase: " + text + " </s>"
encoding = tokenizer.encode_plus(text, padding="longest", return_tensors="pt")
input_ids = encoding["input_ids"].to(device)
attention_masks = encoding["attention_mask"].to(device)
# Generate (Rewrite)
outputs = model.generate(
input_ids=input_ids,
attention_mask=attention_masks,
max_length=256,
do_sample=True, # Adds creativity
top_k=120,
top_p=0.95,
early_stopping=True,
num_return_sequences=1
)
result = tokenizer.decode(outputs[0], skip_special_tokens=True, clean_up_tokenization_spaces=True)
return result
except Exception as e:
return f"Error: {str(e)}"
# 2. INTERFACE
with gr.Blocks() as demo:
with gr.Row():
input_box = gr.Textbox(label="Original Text", lines=5, placeholder="Paste text here...")
output_box = gr.Textbox(label="Rewritten Text", lines=5)
btn = gr.Button("Rewrite Text")
btn.click(fn=paraphrase_text, inputs=input_box, outputs=output_box, api_name="predict")
demo.launch()