Update app.py

app.py

@@ -6,24 +6,22 @@ import torch
 model_name = "haykgrigorian/TimeCapsuleLLM-v2-llama-1.2B"
 
 # 2. Load the Model and Tokenizer
-# We use device_map="auto" to use available CPU RAM efficiently
 print("Loading model... this usually takes 1-2 minutes on first run.")
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # 3. Define the Generate Function
 def generate_text(prompt, max_tokens=100, temperature=0.7):
-    # Format inputs
-    inputs = tokenizer(prompt, return_tensors="pt")
+    # FIX: We added return_token_type_ids=False to stop the error
+    inputs = tokenizer(prompt, return_tensors="pt", return_token_type_ids=False)
     
     # Generate
-    # We disable gradients to save memory and speed up inference
     with torch.no_grad():
         outputs = model.generate(
             **inputs, 
             max_new_tokens=int(max_tokens),
             temperature=float(temperature),
-            do_sample=True, # Allows for creativity/temperature
+            do_sample=True,
             pad_token_id=tokenizer.eos_token_id
         )
     
@@ -31,7 +29,6 @@ def generate_text(prompt, max_tokens=100, temperature=0.7):
     return tokenizer.decode(outputs[0], skip_special_tokens=True)
 
 # 4. Launch the Gradio Interface
-# This creates the UI and the API endpoint automatically
 iface = gr.Interface(
     fn=generate_text,
     inputs=[