Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import torch | |
| import torch.nn.functional as F | |
| import tiktoken | |
| from model import GPT, GPTConfig # We'll create this file next | |
| from huggingface_hub import hf_hub_download | |
| # Initialize model and load weights | |
| def load_model(): | |
| model = GPT(GPTConfig( | |
| block_size=128, | |
| vocab_size=50257, | |
| n_layer=12, | |
| n_head=12, | |
| n_embd=768 | |
| )) | |
| # Download the model from the Hub | |
| model_path = hf_hub_download( | |
| repo_id="amlanr66/gpt2TransformerCustomeTrained", # Replace with your model repository | |
| filename="best_model.pt" | |
| ) | |
| checkpoint = torch.load(model_path, map_location=torch.device('cpu')) | |
| model.load_state_dict(checkpoint['model_state_dict']) | |
| model.eval() | |
| return model | |
| # Text generation function | |
| def generate_text(prompt, max_length=50, temperature=0.8, top_k=50): | |
| # Encode the prompt | |
| enc = tiktoken.get_encoding('gpt2') | |
| input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0) # Add batch dimension | |
| # Generate text | |
| with torch.no_grad(): | |
| for _ in range(max_length): | |
| # Get the last block_size tokens | |
| input_ids_cond = input_ids[:, -model.config.block_size:] | |
| # Forward pass | |
| logits = model(input_ids_cond)[0] | |
| logits = logits[:, -1, :] / temperature | |
| # Apply top-k sampling | |
| if top_k is not None: | |
| v, _ = torch.topk(logits, min(top_k, logits.size(-1))) | |
| logits[logits < v[:, [-1]]] = float('-inf') | |
| # Sample from the distribution | |
| probs = F.softmax(logits, dim=-1) | |
| next_token = torch.multinomial(probs, num_samples=1) | |
| # Append the new token | |
| input_ids = torch.cat((input_ids, next_token), dim=1) | |
| # Decode the generated text | |
| generated_text = enc.decode(input_ids[0].tolist()) | |
| return generated_text | |
| # Load the model globally | |
| model = load_model() | |
| # Create Gradio interface | |
| def predict(text, length, temp): | |
| try: | |
| output = generate_text( | |
| prompt=text, | |
| max_length=int(length), | |
| temperature=float(temp) | |
| ) | |
| return output | |
| except Exception as e: | |
| return f"Error: {str(e)}" | |
| # Define the interface | |
| iface = gr.Interface( | |
| fn=predict, | |
| inputs=[ | |
| gr.Textbox(label="Input Text", placeholder="Enter your prompt here..."), | |
| gr.Slider(minimum=10, maximum=200, value=50, step=1, label="Maximum Length"), | |
| gr.Slider(minimum=0.1, maximum=2.0, value=0.8, step=0.1, label="Temperature") | |
| ], | |
| outputs=gr.Textbox(label="Generated Text"), | |
| title="GPT Text Generation", | |
| description="Enter a prompt and the model will continue the text.", | |
| examples=[ | |
| ["The quick brown fox", 50, 0.8], | |
| ["Once upon a time", 100, 0.9], | |
| ["In a world where", 75, 0.7] | |
| ] | |
| ) | |
| # Launch the app | |
| if __name__ == "__main__": | |
| iface.launch() |