Create app.py

app.py

@@ -0,0 +1,60 @@
+import gradio as gr
+import torch
+from mingru_lm import MinGRU_LM
+
+
+model = MinGRU_LM(dim=512, num_tokens=256, num_layers=6)
+pt_model = "model/best_model.pt"
+checkpoint = torch.load(pt_model)
+model.load_state_dict(checkpoint['model_state_dict'])
+
+# Move model to GPU if available
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+model = model.to(device)
+
+def decode_tokens(tokens):
+    return ''.join([chr(token) for token in tokens if token >= 32 and token < 256])  # ASCII-safe decoding
+
+def tokenize_text(text):
+    return [ord(char) for char in text if ord(char) < 256]  # ASCII-safe tokenization
+
+def generate_text(start_text, max_length, temperature):
+    model.eval()
+    
+    tokens = tokenize_text(start_text)
+    input_tensor = torch.tensor(tokens, dtype=torch.long).unsqueeze(0).to(device)  # Ensure long tensor
+    
+    generated_tokens = tokens.copy()
+    
+    with torch.no_grad():
+        for _ in range(max_length):
+            _, logits = model(input_tensor, labels=None)
+            
+            last_token_logits = logits[0, -1, :] / temperature
+            probs = torch.softmax(last_token_logits, dim=-1)
+            next_token = torch.multinomial(probs, num_samples=1).item()
+            
+            # Only append if it's within the 256-character ASCII range
+            if next_token < 256:  
+                generated_tokens.append(next_token)
+                input_tensor = torch.cat([input_tensor, torch.tensor([[next_token]], device=device)], dim=1)
+            else:
+                continue  # Skip tokens outside ASCII range
+        
+    return decode_tokens(generated_tokens)
+
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_text,
+    inputs=[
+        gr.Textbox(lines=3, label="Enter your prompt", value="Once upon a time"),
+        gr.Slider(minimum=10, maximum=500, value=200, step=1, label="Max Length"),
+        gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
+    ],
+    outputs=gr.Textbox(lines=10, label="Generated Text"),
+    title="Text Generation with MinGRU_LM",
+    description="Enter a prompt and adjust parameters to generate text using the MinGRU_LM model."
+)
+
+if __name__ == "__main__":
+    iface.launch()