Upload directory

inference.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+from safetensors.torch import load_file
+from transformers import BertTokenizer
+
+class Gemma3ForConditionalGeneration(nn.Module):
+    def __init__(self, vocab_size, embedding_dim=1344, num_heads=64, num_layers=48):
+        super(Gemma3ForConditionalGeneration, self).__init__()
+        self.token_embeddings = nn.Embedding(vocab_size, embedding_dim)
+        self.transformer_layers = nn.ModuleList([
+            nn.TransformerEncoderLayer(d_model=embedding_dim, nhead=num_heads) for _ in range(num_layers)
+        ])
+        self.output_layer = nn.Linear(embedding_dim, vocab_size)
+
+    def forward(self, input_ids):
+        text_embeddings = self.token_embeddings(input_ids)
+        for layer in self.transformer_layers:
+            text_embeddings = layer(text_embeddings)
+        return self.output_layer(text_embeddings)
+
+def load_model(model_path, vocab_size):
+    model = Gemma3ForConditionalGeneration(vocab_size=vocab_size)
+    model_weights = load_file(model_path)
+    model.load_state_dict(model_weights, strict=False)
+    model.eval()
+    return model
+
+def generate_text(model, tokenizer, prompt, max_length=50, temperature=1.0):
+    input_ids = tokenizer(prompt, return_tensors='pt').input_ids
+    generated_ids = input_ids
+
+    for _ in range(max_length):
+        with torch.no_grad():
+            outputs = model(generated_ids)
+            next_token_logits = outputs[:, -1, :]  # Get the logits for the last token
+            
+            # Apply temperature
+            next_token_logits = next_token_logits / temperature
+            
+            # Use softmax to get probabilities
+            probabilities = torch.softmax(next_token_logits, dim=-1)
+            
+            # Sample from the distribution
+            next_token = torch.multinomial(probabilities, num_samples=1)  # Sample a token
+            
+            generated_ids = torch.cat((generated_ids, next_token.unsqueeze(0)), dim=1)  # Append the predicted token
+
+    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
+    return generated_text
+
+if __name__ == "__main__":
+    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+    vocab_size = 262208 // 4
+    model_path = './model.safetensors'  # Replace with your model path
+    model = load_model(model_path, vocab_size)
+
+    prompt = "What is the capital of France?"
+    
+    # Generate text based on the prompt with a specified temperature
+    generated_text = generate_text(model, tokenizer, prompt, temperature=0.7)  # Adjust temperature as needed
+    print("Generated Text:", generated_text)