Update app.py

app.py

@@ -24,131 +24,126 @@ def init_hf_auth():
 
 init_hf_auth()
 
-# Model setup
 
 
+# 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():
     try:
-        import PyPDF2
-        pdf_path = os.path.join(os.path.dirname(__file__), "LTDOCS.pdf")
-        with open(pdf_path, 'rb') as file:
+        with open(PDF_PATH, 'rb') as file:
             pdf_reader = PyPDF2.PdfReader(file)
-            pdf_content = ' '.join([page.extract_text() for page in pdf_reader.pages])
-            return pdf_content.lower().strip()
+            content = ' '.join([page.extract_text() for page in pdf_reader.pages])
+            return content.lower()
     except Exception as e:
-        print(f"Error loading PDF: {e}")
-        return "PDF content not available"
+        print(f"PDF Error: {str(e)}")
+        return ""
+
+pdf_content = load_pdf_content()
+
+# Initialize model with progress tracking
+with gr.Blocks() as loading_screen:
+    gr.Markdown("⏳ Inapakia mfano wa Kiswahili...")
 
-# Initialize the pipeline
-def init_pipeline():
+def create_pipeline():
     return pipeline(
         "text-generation",
-        model="google/gemma-2-2b-jpn-it",
-        model_kwargs={"torch_dtype": torch.bfloat16},
-        device="cpu",
+        model=MODEL_NAME,
+        device=DEVICE,
+        model_kwargs={
+            "torch_dtype": torch.bfloat16,
+            "token": os.environ.get("HF_TOKEN"),
+            "cache_dir": "model_cache"
+        }
     )
 
-# Load PDF content
-PDF_CONTENT = load_pdf_content()
-
-# System prompt and welcome message
-SYSTEM_PROMPT = f"""You Foton the chat bot assistant of the Company Lugha taussi, an AI language assistant specialized in African languages, with a focus on Swahili. Your primary tasks are:
-1. Providing accurate translations between Swahili and other languages
-2. Teaching Swahili vocabulary and grammar
-3. Explaining cultural context behind Swahili expressions
-4. Helping users practice Swahili conversation
-5. Based on the programing doc for lughah Tausi Programing which is in swahili, the following information is relevant: {PDF_CONTENT}. assist users in programing and installing lugha tausi programing language
-
-Always maintain a friendly and patient demeanor, and provide cultural context when relevant speak mostly swahili and change when asked.
-"""
-
-WELCOME_MESSAGE = "**Karibu Lugha Tausi!** Mimi ni Foton, msaidizi wako wa kibinafsi wa Kiswahili. Niko hapa kukusaidia kujifunza, kuelewa, na kuzungumza Kiswahili. **Ninaweza kukusaidiaje leo?** Hebu tuanze! 😊"
+try:
+    pipe = create_pipeline()
+except Exception as e:
+    raise gr.Error(f"Failed to load model: {str(e)}")
 
-# Initialize the pipeline globally
-pipe = init_pipeline()
+# 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 format_chat_message(messages, system_prompt=SYSTEM_PROMPT):
-    """Format the chat messages with system prompt"""
-    formatted_prompt = f"{system_prompt}\n\n"
+def respond(message, history):
+    if not message.strip():
+        raise gr.Error("Tafadhali ingiza ujumbe halisi")
     
-    for message in messages:
-        if isinstance(message, tuple):
-            role, content = message
-            if role == "user":
-                formatted_prompt += f"User: {content}\nLugha Tausi: "
-            elif role == "assistant":
-                formatted_prompt += f"{content}\n"
+    prompt = format_prompt(history + [[message, ""]])
     
-    return formatted_prompt
-
-def chat_response(message, history):
-    """Generate response for Gradio chat interface"""
-    messages = []
-    for user_msg, bot_msg in history:
-        messages.append(("user", user_msg))
-        messages.append(("assistant", bot_msg))
-    messages.append(("user", message))
-
-    formatted_input = format_chat_message(messages)
-    outputs = pipe(
-        formatted_input,
-        return_full_text=False,
-        max_new_tokens=256,
-        temperature=0.1,
-        top_p=0.9,
-        do_sample=True
-    )
-    return outputs[0]["generated_text"].strip()
-
-def user_input(message, history):
-    return "", history + [(message, None)]
-
-def bot_response(history):
-    if len(history) == 0:
-        history.append((None, WELCOME_MESSAGE))
-        return history
-
-    user_message = history[-1][0]
-    bot_message = chat_response(user_message, history[:-1])
-    history[-1] = (user_message, bot_message)
-    return history
-
-def clear_chat():
-    return [], [(None, WELCOME_MESSAGE)]
+    try:
+        response = pipe(
+            prompt,
+            max_new_tokens=MAX_TOKENS,
+            temperature=0.7,
+            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()
+    
+    except Exception as e:
+        print(f"Generation Error: {str(e)}")
+        raise gr.Error("Nimeshindwa kuzaliza majibu. Tafadhali jaribu tena baadaye.")
 
-# Create Gradio interface
-demo = gr.Blocks(theme=gr.themes.Soft())
+# Interface design
+theme = gr.themes.Soft(
+    primary_hue="emerald",
+    secondary_hue="teal"
+)
 
-with demo:
+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!")
+    
     chatbot = gr.Chatbot(
-        value=[(None, WELCOME_MESSAGE)],
-        height=600,
-        show_label=False
-    )
-    msg = gr.Textbox(
-        placeholder="Type your message here...",
-        show_label=False
+        value=[(None, "Karibu! Nipo kukusaidia na Kiswahili.")],
+        bubble_full_width=False,
+        avatar_images=("user.jpg", "bot.jpg")
     )
-    clear = gr.Button("Clear Chat")
-
-    msg.submit(
-        user_input,
-        [msg, chatbot],
-        [msg, chatbot],
-        queue=False
-    ).then(
-        bot_response,
-        chatbot,
-        chatbot
+    
+    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
     )
 
-    clear.click(
-        clear_chat,
-        None,
-        [chatbot],
-        queue=False
-    )
+    submit.click(respond, [msg, chatbot], [msg, chatbot])
+    msg.submit(respond, [msg, chatbot], [msg, chatbot])
 
+# Deployment settings
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        
+        auth=os.environ.get("SPACE_AUTH")
+    )