Update app.py

app.py

@@ -1,43 +1,48 @@
+
 import torch
 import gradio as gr
-from transformers import GPT2Tokenizer, GPT2LMHeadModel
-
-# Load the tokenizer and the model
-tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
-model = GPT2LMHeadModel.from_pretrained('gpt2')
+from model import GPT, GPTConfig  # Assuming your model code is in a file named model.py
+import tiktoken
 
-# Load the best model weights
-model.load_state_dict(torch.load('GPT_model.pth', map_location=torch.device('cpu')))
+# Load the trained model
+def load_model(model_path):
+    config = GPTConfig()  # Adjust this if you've changed the default config
+    model = GPT(config)
+    model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))
+    model.eval()
+    return model
 
-# Set the model to evaluation mode
-model.eval()
+model = load_model('GPT_model.pth')  # Replace with the actual path to your .pth file
+enc = tiktoken.get_encoding('gpt2')
 
-# Define the text generation function
-def generate_text(prompt, max_length=50, num_return_sequences=1):
-    inputs = tokenizer(prompt, return_tensors='pt')
-    outputs = model.generate(
-        inputs.input_ids,
-        max_length=max_length,
-        num_return_sequences=num_return_sequences,
-        do_sample=True,
-        top_k=50,
-        top_p=0.95,
-        temperature=1.0
-    )
-    return [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+def generate_text(prompt, max_length=100, temperature=0.7):
+    input_ids = torch.tensor(enc.encode(prompt)).unsqueeze(0)
+    
+    with torch.no_grad():
+        for _ in range(max_length):
+            outputs = model(input_ids)
+            next_token_logits = outputs[0][:, -1, :] / temperature
+            next_token = torch.multinomial(torch.softmax(next_token_logits, dim=-1), num_samples=1)
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            
+            if next_token.item() == enc.encode('\n')[0]:
+                break
+    
+    generated_text = enc.decode(input_ids[0].tolist())
+    return generated_text
 
-# Define the Gradio interface
-interface = gr.Interface(
+# Gradio interface
+iface = gr.Interface(
     fn=generate_text,
     inputs=[
-        gr.inputs.Textbox(lines=2, placeholder="Enter your prompt here..."),
-        gr.inputs.Slider(minimum=10, maximum=200, default=50, label="Max Length"),
-        gr.inputs.Slider(minimum=1, maximum=5, default=1, label="Number of Sequences")
+        gr.Textbox(label="Prompt", placeholder="Enter your prompt here..."),
+        gr.Slider(minimum=10, maximum=200, value=100, step=1, label="Max Length"),
+        gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
     ],
-    outputs=gr.outputs.Textbox(),
+    outputs=gr.Textbox(label="Generated Text"),
     title="GPT-2 Text Generator",
-    description="Enter a prompt to generate text using GPT-2.",
+    description="Enter a prompt and generate text using a fine-tuned GPT-2 model."
 )
 
-# Launch the Gradio interface
-interface.launch()
+# Launch the app
+iface.launch()