Update app.py

app.py

@@ -1,89 +1,96 @@
-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
-
-# 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
-    ))
-    
-    checkpoint = torch.load('best_model.pt', 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__":
+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()