Upload app.py

app.py

@@ -1,8 +1,9 @@
 import gradio as gr
 import torch
 import torch.nn as nn
+import time
 
-# Define your custom model class with detailed layer structures
+# Define the custom model class with detailed layer structures
 class Head(nn.Module):
     def __init__(self, head_size):
         super().__init__()
@@ -103,7 +104,6 @@ def load_model():
 model = load_model()
 
 # Define a comprehensive character set based on training data
-# Convert all input to lowercase if the model is trained on lowercase data
 chars = sorted(list(set("abcdefghijklmnopqrstuvwxyz0123456789 .,!?-:;'\"\n")))
 stoi = {ch: i for i, ch in enumerate(chars)}
 itos = {i: ch for i, ch in enumerate(chars)}
@@ -113,6 +113,7 @@ decode = lambda l: ''.join([itos[i] for i in l])
 # Function to generate text using the model
 def generate_text(prompt):
     try:
+        start_time = time.time()
         print(f"Received prompt: {prompt}")
         encoded_prompt = encode(prompt)
         
@@ -128,11 +129,12 @@ def generate_text(prompt):
         print(f"Encoded prompt: {context}")
         
         with torch.no_grad():
-            generated = model.generate(context, max_new_tokens=250)  # Adjust as needed
+            generated = model.generate(context, max_new_tokens=20)  # Reduced tokens to speed up
             print(f"Generated tensor: {generated}")
         
         result = decode(generated[0].tolist())
         print(f"Decoded result: {result}")
+        print(f"Processing time: {time.time() - start_time:.2f}s")
         return result
     except Exception as e:
         print(f"Error during generation: {e}")
@@ -144,7 +146,8 @@ interface = gr.Interface(
     inputs=gr.Textbox(lines=2, placeholder="Enter a location or prompt..."),
     outputs="text",
     title="Triptuner Model",
-    description="Generate itineraries for locations in Sri Lanka's Central Province."
+    description="Generate itineraries for locations in Sri Lanka's Central Province.",
+    theme="compact",  # Add a theme for better UI appearance
 )
 
 # Launch the interface