Update app.py

app.py

@@ -5,44 +5,74 @@ import torch
 # Load model and tokenizer
 model_name = "0dAI/0dAI-8x7b-0761"
 
-@cache_resource
+@gr.cache_resource
 def load_model():
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float16,
-        device_map="auto",
-        trust_remote_code=True
-    )
-    return tokenizer, model
+    try:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForCausalLM.from_pretrained(
+            model_name,
+            torch_dtype=torch.float16,
+            device_map="auto",
+            trust_remote_code=True,
+            low_cpu_mem_usage=True
+        )
+        return tokenizer, model
+    except Exception as e:
+        print(f"Error loading model: {e}")
+        return None, None
 
 def generate_text(prompt, max_length=512):
     tokenizer, model = load_model()
     
-    inputs = tokenizer(prompt, return_tensors="pt")
-    with torch.no_grad():
-        outputs = model.generate(
-            **inputs,
-            max_length=max_length,
-            num_return_sequences=1,
-            temperature=0.7,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
+    if tokenizer is None or model is None:
+        return "Error: Model failed to load. Please check the logs."
     
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return response
+    try:
+        # Tokenize input
+        inputs = tokenizer(prompt, return_tensors="pt")
+        
+        # Move to same device as model
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=max_length,
+                temperature=0.7,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                repetition_penalty=1.1
+            )
+        
+        # Decode response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        return response[len(prompt):]  # Return only the generated part
+        
+    except Exception as e:
+        return f"Error during generation: {str(e)}"
 
 # Create Gradio interface
-iface = gr.Interface(
+demo = gr.Interface(
     fn=generate_text,
     inputs=[
-        gr.Textbox(label="Input Prompt", lines=3),
-        gr.Slider(100, 1024, value=512, label="Max Length")
+        gr.Textbox(
+            label="Input Prompt",
+            lines=3,
+            placeholder="Enter your prompt here...",
+            value="Hello, how are you?"
+        ),
+        gr.Slider(50, 1024, value=256, label="Max New Tokens")
     ],
-    outputs=gr.Textbox(label="Generated Text", lines=5),
+    outputs=gr.Textbox(label="Generated Text", lines=8),
     title="0dAI 8x7B Model Demo",
-    description="Interactive demo for the 0dAI/0dAI-8x7b-0761 model"
+    description="Interactive demo for the 0dAI/0dAI-8x7b-0761 model. This is a large model, so initial loading may take a few minutes.",
+    examples=[
+        ["Explain quantum computing in simple terms."],
+        ["Write a short story about a robot learning to paint."],
+        ["What are the benefits of renewable energy?"]
+    ]
 )
 
-iface.launch()
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)