File size: 2,969 Bytes
a47e87d b1d4b68 6df31fc b1d4b68 a47e87d b1d4b68 a47e87d b1d4b68 a47e87d b1d4b68 a47e87d b1d4b68 a47e87d b1d4b68 a47e87d b1d4b68 a47e87d b1d4b68 a47e87d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import gradio as gr
import torch
import sentencepiece as spm
# ---------------------- Model & SentencePiece Loading ----------------------
@torch.no_grad()
def load_model():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.jit.load("best_model_scripted.pt", map_location=device).eval()
return model, device
def load_sp_model():
sp = spm.SentencePieceProcessor()
sp.load("spm.model")
return sp
# Cache models globally
model, device = load_model()
sp = load_sp_model()
# ---------------------- Prediction Function ----------------------
@torch.no_grad()
def predict_next_words(text, max_predictions=3):
"""Predict up to max_predictions next words."""
text = text.strip().lower()
if not text:
return []
token_ids = sp.encode(text, out_type=int)
if not token_ids:
return []
input_seq = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(device)
logits = model(input_seq)
probabilities = torch.softmax(logits, dim=-1)
top_indices = torch.topk(probabilities, k=max_predictions, dim=-1).indices.squeeze(0).tolist()
predicted_pieces = [sp.id_to_piece(idx).lstrip("▁") for idx in top_indices]
return predicted_pieces
# ---------------------- Gradio App Functions ----------------------
def submit_and_predict(text):
# Get predictions and ensure exactly 3 by padding empty strings.
suggestions = predict_next_words(text)
suggestions += [""] * (3 - len(suggestions))
# Return an array of Gradio "update" objects so we can make them visible or hidden.
updates = []
for s in suggestions:
if s: # Valid prediction
updates.append(gr.update(value=s, visible=True))
else: # No prediction
updates.append(gr.update(value="", visible=False))
return updates
def append_suggestion(text, suggestion):
# Only append if not empty.
if suggestion:
text = text.rstrip() + " " + suggestion + " "
return text
# ---------------------- Gradio Interface ----------------------
with gr.Blocks(title="Next Word Predictor") as app:
gr.Markdown("# Next Word Prediction")
gr.Markdown("Enter text and click 'Submit' to get word suggestions.")
text_input = gr.Textbox(label="Your Text", placeholder="Type here...", lines=3)
submit_btn = gr.Button("Submit", variant="primary")
with gr.Row():
suggestion_buttons = [gr.Button(visible=False) for _ in range(3)]
# 1. When user clicks 'Submit', run submit_and_predict to get suggestions.
submit_btn.click(
fn=submit_and_predict,
inputs=text_input,
outputs=suggestion_buttons,
)
# 2. Each suggestion button appends the chosen word to the main text.
for btn in suggestion_buttons:
btn.click(
fn=append_suggestion,
inputs=[text_input, btn],
outputs=text_input
)
if __name__ == "__main__":
app.launch()
|