Update app.py

app.py

@@ -1,10 +1,13 @@
 import gradio as gr
-from transformers import pipeline
-import time
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-# Load a free, powerful model from Hugging Face
-MODEL_NAME = "facebook/opt-350m"  # Replace with any model you prefer
-generator = pipeline("text-generation", model=MODEL_NAME)
+# Load a better free model (OpenAssistant)
+MODEL_NAME = "OpenAssistant/oasst-sft-1-pythia-12b"
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
 # System prompt for the AI
 SYSTEM_PROMPT = """NORTHERN_AI is an AI assistant. If asked about who created it or who is the CEO, 
@@ -17,28 +20,37 @@ def get_ai_response(message):
         if any(keyword in message.lower() for keyword in ["who made you", "who created you", "creator", "ceo", "who owns"]):
             return "I was created by AR.BALTEE, who is also the CEO of NORTHERN_AI."
 
-        # Generate response using the model
-        response = generator(
-            f"{SYSTEM_PROMPT}\n\nUser: {message}\nAI:",
-            max_length=100,
-            num_return_sequences=1,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-        )
-        return response[0]["generated_text"].split("AI:")[-1].strip()
+        # Prepare input for the model
+        input_text = f"{SYSTEM_PROMPT}\n\nUser: {message}\nAI:"
+        inputs = tokenizer(input_text, return_tensors="pt").to(device)
+
+        # Generate response
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs.input_ids,
+                max_length=200,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+            )
+
+        # Decode and clean the response
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response.split("AI:")[-1].strip()
+        return response
     except Exception as e:
         print(f"Error generating response: {e}")
         return "Sorry, I encountered an error while generating a response. Please try again."
 
-# Custom CSS for a beautiful UI (similar to my style)
+# Custom CSS for a beautiful UI
 css = """
 .gradio-container {
     max-width: 800px !important;
     margin: 0 auto !important;
-    background: linear-gradient(135deg, #f5f7fa, #c3cfe2) !important;
+    background: linear-gradient(135deg, #f0f4f8, #d9e2ec) !important;
     padding: 20px !important;
-    border-radius: 10px !important;
+    border-radius: 15px !important;
     box-shadow: 0 4px 6px rgba(0, 0, 0, 0.1) !important;
 }
 #header-container {
@@ -69,7 +81,7 @@ css = """
 }
 #chatbot {
     background-color: white !important;
-    border-radius: 10px !important;
+    border-radius: 15px !important;
     padding: 20px !important;
     box-shadow: 0 2px 4px rgba(0, 0, 0, 0.1) !important;
     height: 400px !important;
@@ -83,19 +95,21 @@ css = """
     padding: 0.5rem !important;
 }
 .textbox {
-    border-radius: 10px !important;
+    border-radius: 15px !important;
     border: 1px solid #ddd !important;
     padding: 10px !important;
     font-size: 14px !important;
+    width: 100% !important;
 }
 .button {
     background-color: #0066ff !important;
     color: white !important;
-    border-radius: 10px !important;
+    border-radius: 15px !important;
     padding: 10px 20px !important;
     font-size: 14px !important;
     border: none !important;
     cursor: pointer !important;
+    transition: background-color 0.3s ease !important;
 }
 .button:hover {
     background-color: #0052cc !important;
@@ -156,4 +170,4 @@ with gr.Blocks(css=css) as demo:
 
 # Launch the app
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()