Add updated filed

app.py

@@ -1,27 +1,28 @@
 import streamlit as st
-from transformers import GPT2LMHeadModel, AutoTokenizer, AutoConfig
+import torch
+from transformers import AutoModelForCausalLM
 
-# Load the configuration
-config = AutoConfig.from_pretrained("praty7717/Odeyssey")
+# Define your custom config class
+class MyCustomConfig(PretrainedConfig):
+    model_type = "gpt"
 
-# Load the model
-model = GPT2LMHeadModel.from_pretrained("praty7717/Odeyssey", config=config)
+    def __init__(self, vocab_size, n_embd, n_layer, n_head, block_size, **kwargs):
+        super().__init__(vocab_size=vocab_size, n_embd=n_embd, n_layer=n_layer, n_head=n_head, block_size=block_size, **kwargs)
 
-# Load the tokenizer
-tokenizer = AutoTokenizer.from_pretrained("praty7717/Odeyssey")
+# Load the model and configuration
+config = MyCustomConfig.from_pretrained("praty7717/Odeyssey")
+model = AutoModelForCausalLM.from_pretrained("praty7717/Odeyssey", config=config)
 
-# Define the text generation function
+# Function to generate text
 def generate_text(prompt, max_length=100):
-    input_ids = tokenizer.encode(prompt, return_tensors='pt')
-    output = model.generate(input_ids, max_length=max_length, num_return_sequences=1)
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    inputs = tokenizer(prompt, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=max_length)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # Streamlit UI
 st.title("Text Generation with Custom GPT Model")
-
-# User input for the prompt
-user_input = st.text_input("Enter a prompt:", "Once upon a time")
-
+start_prompt = st.text_area("Enter your prompt here:", "Once upon a time")
 if st.button("Generate"):
-    generated_text = generate_text(user_input, max_length=100)
+    generated_text = generate_text(start_prompt, max_length=100)
     st.write(generated_text)