Upload app.py

app.py

@@ -1,18 +1,12 @@
 import torch
 import gradio as gr
-from model import CustomLLM
-from transformers import GPT2Tokenizer
+from model import CustomLLM, CustomConfig
+from transformers import AutoTokenizer
 
 class ModelLoader:
     def __init__(self):
         # Load config
-        self.config = {
-            "vocab_size": 50257,  # Update with your actual values
-            "hidden_size": 768,
-            "num_hidden_layers": 12,
-            "rms_norm_eps": 1e-6
-        }
-        
+        self.config = CustomConfig()
         # Instantiate model
         self.model = CustomLLM(self.config)
         
@@ -22,7 +16,7 @@ class ModelLoader:
         self.model.eval()
         
         # Load tokenizer
-        self.tokenizer = GPT2Tokenizer.from_pretrained('tokenizer/')
+        self.tokenizer = AutoTokenizer.from_pretrained("HuggingFaceTB/cosmo2-tokenizer")
         self.tokenizer.pad_token = self.tokenizer.eos_token
 
     def generate(self, prompt, max_new_tokens=100, temperature=0.9, top_k=50, top_p=0.95):
@@ -36,7 +30,7 @@ class ModelLoader:
                 temperature=temperature,
                 top_k=top_k,
                 top_p=top_p,
-                eos_token_id=self.tokenizer.eos_token_id,
+                eos_token_id=None,
                 pad_token_id=self.tokenizer.pad_token_id
             )
         
@@ -60,4 +54,4 @@ interface = gr.Interface(
     description="Generate text using your custom-trained LLM"
 )
 
-interface.launch()
+interface.launch(share=True)