File size: 2,338 Bytes
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 |
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):
if not text.strip():
return []
token_ids = sp.encode(text.strip(), out_type=int)
if len(token_ids) == 0:
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):
suggestions = predict_next_words(text)
suggestions += [""] * (3 - len(suggestions)) # Ensure 3 buttons always
return suggestions
def append_suggestion(text, suggestion):
if suggestion:
text += 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)]
submit_btn.click(
fn=submit_and_predict,
inputs=text_input,
outputs=suggestion_buttons,
)
for btn in suggestion_buttons:
btn.click(
fn=append_suggestion,
inputs=[text_input, btn],
outputs=text_input
)
if __name__ == "__main__":
app.launch()
|