szili2011's picture
Update app.py
0f412e0 verified
import torch
import torch.nn as nn
import json
import gradio as gr
# --- Step 1: Create a "Smart" Vocabulary Loader ---
# This function will load the vocabularies and automatically fix any format mismatches.
def load_vocabularies():
"""
Loads vocabularies and intelligently determines the correct format,
preventing crashes due to misnamed files.
"""
with open('char_to_int.json', 'r', encoding='utf-8') as f:
char_to_int_map = json.load(f)
# Load the file the user has named 'int_to_lang.json'.
with open('int_to_lang.json', 'r', encoding='utf-8') as f:
language_vocab = json.load(f)
# Get the first key to check the format (e.g., is it "0" or "C#")
first_key = next(iter(language_vocab))
int_to_lang_map = {}
try:
# Try to convert the first key to an integer.
int(first_key)
# If this SUCCEEDS, the file is in the correct {"0": "Language"} format.
print("[INFO] Detected int->lang format. Loading directly.")
int_to_lang_map = {int(k): v for k, v in language_vocab.items()}
except ValueError:
# If this FAILS, the file is in the {"Language": 0} format.
# We must reverse it to create the correct int->lang map.
print("[INFO] Detected lang->int format. Reversing dictionary to fix.")
int_to_lang_map = {v: k for k, v in language_vocab.items()}
return char_to_int_map, int_to_lang_map
# Load the vocabularies using our smart function
char_to_int, int_to_lang = load_vocabularies()
# --- Step 2: Re-define the Model Architecture ---
# This MUST be the exact same architecture as the one you trained.
class CodeClassifierRNN(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim, output_dim, n_layers, bidirectional, dropout, pad_idx):
super().__init__()
self.embedding = nn.Embedding(vocab_size, embedding_dim, padding_idx=pad_idx)
self.lstm = nn.LSTM(embedding_dim, hidden_dim, num_layers=n_layers, bidirectional=bidirectional, dropout=dropout if n_layers > 1 else 0, batch_first=True)
self.dropout = nn.Dropout(dropout)
self.fc = nn.Linear(hidden_dim * 2, output_dim) # * 2 for bidirectional
def forward(self, text):
embedded = self.embedding(text)
_, (hidden, _) = self.lstm(embedded)
hidden = torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim=1)
hidden = self.dropout(hidden)
output = self.fc(hidden)
return output
# --- Step 3: Instantiate the model and load the trained weights ---
PAD_IDX = char_to_int['<PAD>']
VOCAB_SIZE = len(char_to_int)
EMBEDDING_DIM = 128
HIDDEN_DIM = 192
OUTPUT_DIM = len(int_to_lang)
N_LAYERS = 2
BIDIRECTIONAL = True
DROPOUT = 0.5
model = CodeClassifierRNN(VOCAB_SIZE, EMBEDDING_DIM, HIDDEN_DIM, OUTPUT_DIM, N_LAYERS, BIDIRECTIONAL, DROPOUT, PAD_IDX)
model.load_state_dict(torch.load('polyglot_classifier.pt', map_location='cpu'))
model.eval()
# --- Step 4: Create the prediction function ---
def classify_code(code_snippet):
if not code_snippet or not code_snippet.strip():
return {}
indexed = [char_to_int.get(c, char_to_int['<UNK>']) for c in code_snippet]
tensor = torch.LongTensor(indexed).unsqueeze(0)
with torch.no_grad():
prediction = model(tensor)
probabilities = torch.softmax(prediction, dim=1)
top5_probs, top5_indices = torch.topk(probabilities, 5)
# This lookup will now work regardless of the original file format.
confidences = {int_to_lang[idx.item()]: prob.item() for idx, prob in zip(top5_indices[0], top5_probs[0])}
return confidences
# --- Step 5: Create and launch the Gradio Interface ---
iface = gr.Interface(
fn=classify_code,
inputs=gr.Code(language=None, label="Code Snippet"),
outputs=gr.Label(num_top_classes=5, label="Predicted Language"),
title="Polyglot Code Classifier",
description="Enter a code snippet to see which programming language the AI thinks it is. This model was trained from scratch on a custom dataset.",
examples=[
["def hello_world():\n print('Hello from Python!')"],
["function greet() {\n console.log('Hello from JavaScript!');\n}"],
["public class Main {\n public static void main(String[] args) {\n System.out.println(\"Hello, Java!\");\n }\n}"]
]
)
iface.launch()