Upload app.py with huggingface_hub

app.py

@@ -1,26 +1,49 @@
+
 import gradio as gr
 import torch
-from transformers import GPT2LMHeadModel, GPT2Tokenizer
+import tiktoken
+from supplementary import GPTModel, generate_text_simple
+
+# Load model configuration
+GPT_CONFIG_124M = {
+    "vocab_size": 50257,
+    "context_length": 1024,
+    "emb_dim": 768,
+    "n_heads": 12,
+    "n_layers": 12,
+    "drop_rate": 0.1,
+    "qkv_bias": False
+}
+
+# Initialize model
+model = GPTModel(GPT_CONFIG_124M)
+
+# Load the trained weights
+model.load_state_dict(torch.load("my_gpt_model.pth", map_location=torch.device('cpu')))
+model.eval()
 
-# Load the model and tokenizer
-model = GPT2LMHeadModel.from_pretrained("samwell/SamGPT")
-tokenizer = GPT2Tokenizer.from_pretrained("gpt2")  # Assuming you used the GPT-2 tokenizer
+tokenizer = tiktoken.get_encoding("gpt2")
 
-def generate_text(prompt, max_length):
-    input_ids = tokenizer.encode(prompt, return_tensors="pt")
-    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1, no_repeat_ngram_size=2)
-    generated_text = tokenizer.decode(output[0], skip_special_tokens=True)
-    return generated_text
+def generate(prompt, max_new_tokens):
+    token_ids = tokenizer.encode(prompt)
+    input_ids = torch.tensor(token_ids).unsqueeze(0)
+    output_ids = generate_text_simple(
+        model=model,
+        idx=input_ids,
+        max_new_tokens=max_new_tokens,
+        context_size=GPT_CONFIG_124M["context_length"]
+    )
+    return tokenizer.decode(output_ids.squeeze(0).tolist())
 
 iface = gr.Interface(
-    fn=generate_text,
+    fn=generate,
     inputs=[
         gr.Textbox(label="Prompt"),
-        gr.Slider(minimum=10, maximum=200, value=50, step=1, label="Max Length")
+        gr.Slider(minimum=1, maximum=100, value=20, step=1, label="Max New Tokens")
     ],
     outputs=gr.Textbox(label="Generated Text"),
     title="SamGPT Text Generation",
-    description="Enter a prompt to generate text with SamGPT."
+    description="Enter a prompt to generate text with the custom language model."
 )
 
-iface.launch()
+iface.launch()