Update app.py

app.py

@@ -8,142 +8,120 @@ from huggingface_hub import login
 # Space configuration
 SPACE_DIR = os.environ.get("HF_HOME", os.getcwd())
 
+# Load and preprocess the PDF content
+try:
+    pdf_path = os.path.join(SPACE_DIR, "LTDOCS.pdf")
+    with open(pdf_path, 'rb') as file:
+        pdf_reader = PyPDF2.PdfReader(file)
+        pdf_content = ' '.join([page.extract_text() for page in pdf_reader.pages])
+        pdf_content = pdf_content.lower().strip()
+except Exception as e:
+    pdf_content = ""
+    print(f"Error loading PDF: {e}")
 
-# Authentication
-def init_hf_auth():
-    if token := os.getenv("HUGGINGFACE_TOKEN"):
-        try:
-            login(token=token, add_to_git_credential=False)
-            print("HF Auth: Success")
-            return True
-        except Exception as e:
-            print(f"HF Auth Error: {e}")
-            return False
-    print("HF Auth: No token found")
-    return False
-
-init_hf_auth()
-
-
-
-# Configuration
-MODEL_NAME = "google/gemma-2-2b-jpn-it"
-PDF_PATH = "LTDOCS.pdf"
-MAX_TOKENS = 512
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-# System prompt template
-SYSTEM_PROMPT = """You are Foton..."""  # Keep your original prompt
-
-# Load PDF content with error handling
-def load_pdf_content():
+def init_huggingface_auth():
+    """Space-friendly authentication"""
+    token = os.getenv("HUGGINGFACE_TOKEN")
+    
+    if not token:
+        print("No HF token found in environment")
+        return False
+    
     try:
-        with open(PDF_PATH, 'rb') as file:
-            pdf_reader = PyPDF2.PdfReader(file)
-            content = ' '.join([page.extract_text() for page in pdf_reader.pages])
-            return content.lower()
+        login(token=token, add_to_git_credential=False)
+        print("HF authentication successful")
+        return True
     except Exception as e:
-        print(f"PDF Error: {str(e)}")
-        return ""
+        print(f"Login error: {e}")
+        return False
 
-pdf_content = load_pdf_content()
+if not init_huggingface_auth():
+    print("Warning: Authentication failed")
 
-# Initialize model with progress tracking
-with gr.Blocks() as loading_screen:
-    gr.Markdown("⏳ Inapakia mfano wa Kiswahili...")
-
-def create_pipeline():
-    return pipeline(
+# Initialize the pipeline with CPU optimization
+try:
+    pipe = pipeline(
         "text-generation",
-        model=MODEL_NAME,
-        device=DEVICE,
+        model="google/gemma-2-2b-jpn-it",
+        device_map="auto",
         model_kwargs={
-            "torch_dtype": torch.bfloat16,
-            "token": os.environ.get("HF_TOKEN"),
-            "cache_dir": "model_cache"
+            "torch_dtype": torch.float16,
+            "low_cpu_mem_usage": True
         }
     )
-
-try:
-    pipe = create_pipeline()
 except Exception as e:
-    raise gr.Error(f"Failed to load model: {str(e)}")
-
-# Chat processing
-def format_prompt(history):
-    full_prompt = SYSTEM_PROMPT.format(pdf_content=pdf_content)
-    for user_msg, bot_msg in history:
-        full_prompt += f"\nUser: {user_msg}\nAssistant: {bot_msg or ''}"
-    return full_prompt
-
-def respond(message, history):
-    if not message.strip():
-        raise gr.Error("Tafadhali ingiza ujumbe halisi")
-    
-    prompt = format_prompt(history + [[message, ""]])
-    
+    print(f"Model loading error: {e}")
+    raise
+
+SYSTEM_PROMPT = f"""You are Foton, Swahili AI assistant. Tasks:
+1. Swahili translations
+2. Teach Swahili vocabulary/grammar
+3. Explain cultural context
+4. Help practice conversations
+5. Programming assistance using: {pdf_content}
+
+Maintain friendly, patient demeanor with cultural context.
+"""
+
+def format_messages(history):
+    """Format chat history with system prompt"""
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
+    for msg in history:
+        if isinstance(msg, dict) and "role" in msg and "content" in msg:
+            messages.append(msg)
+    return messages
+
+def generate_response(history):
+    """Generate AI response"""
     try:
-        response = pipe(
+        messages = format_messages(history)
+        
+        prompt = "\n".join([f"{m['role']}: {m['content']}" for m in messages])
+        
+        output = pipe(
             prompt,
-            max_new_tokens=MAX_TOKENS,
-            temperature=0.7,
+            max_new_tokens=256,
+            temperature=0.1,
             top_p=0.9,
             do_sample=True,
-            truncation=True
-        )[0]['generated_text']
-        
-        # Extract only the new response
-        return response[len(prompt):].split("User:")[0].strip()
-    
+            return_full_text=False
+        )
+        return output[0]["generated_text"].strip()
     except Exception as e:
-        print(f"Generation Error: {str(e)}")
-        raise gr.Error("Nimeshindwa kuzaliza majibu. Tafadhali jaribu tena baadaye.")
-
-# Interface design
-theme = gr.themes.Soft(
-    primary_hue="emerald",
-    secondary_hue="teal"
-)
+        print(f"Generation error: {e}")
+        return "Samahani, nimekutana na tatizo. Tafadhali jaribu tena baadaye."
 
-with gr.Blocks(theme=theme, title="Lugha Tausi") as demo:
-    gr.Markdown("# 🇹🇿 Foton - Msaidizi Wa Kiswahili")
-    
-    with gr.Row():
-        gr.Image("logo.png", width=100, show_label=False)
-        gr.Markdown("### Karibu kwa Huduma ya Lugha Tausi!")
+# Create Gradio interface with Space optimizations
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Lugha Tausi - Swahili Assistant")
     
     chatbot = gr.Chatbot(
-        value=[(None, "Karibu! Nipo kukusaidia na Kiswahili.")],
+        value=[{"role": "assistant", "content": "**Karibu Lugha Tausi!** Mimi ni Foton, msaidizi wako wa Kiswahili. Niko hapa kukusaidia kujifunza na kuzungumza Kiswahili. **Ninaweza kukusaidiaje leo?** 😊"}],
+        height=600,
+        show_label=False,
+        avatar_images=(None, "user.png"),
         bubble_full_width=False,
-        avatar_images=("user.jpg", "bot.jpg")
+        show_share_button=False,
+        type="messages"
     )
     
-    with gr.Row():
-        msg = gr.Textbox(
-            placeholder="Andika hapa...",
-            max_lines=3,
-            autofocus=True
-        )
-        submit = gr.Button("Tuma", variant="primary")
-    
-    examples = gr.Examples(
-        examples=[
-            ["Nisaidie kutafsiri sentensi hii"],
-            ["Ninauliza kuhusu sarufi ya Kiswahili"],
-            ["Eleza maana ya 'Haraka haraka haina baraka'"]
-        ],
-        inputs=msg
-    )
+    msg = gr.Textbox(placeholder="Andika ujumbe wako hapa...", show_label=False)
+    clear = gr.Button("Futa Mazungumzo")
+
+    def respond(message, history):
+        history.append({"role": "user", "content": message})
+        bot_response = generate_response(history)
+        history.append({"role": "assistant", "content": bot_response})
+        return "", history
 
-    submit.click(respond, [msg, chatbot], [msg, chatbot])
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
+    clear.click(lambda: [None, []], None, [msg, chatbot], queue=False)
 
-# Deployment settings
 if __name__ == "__main__":
     demo.launch(
         server_name="0.0.0.0",
         server_port=7860,
-        
-        auth=os.environ.get("SPACE_AUTH")
-    )
-
+        auth=os.getenv("SPACE_AUTH"),
+        show_error=True,
+        share=True  # Explicitly disable sharing
+    )