Create app.py

app.py

@@ -0,0 +1,53 @@
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+# Check for GPU and use CPU if not available (your zero-GPU approach)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print(f"Using device: {device}")
+
+# Load model and tokenizer
+model_id = "google/gemma-2-2b-it"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map="auto",
+    torch_dtype=torch.bfloat16  # Uses less memory
+)
+
+def chat_response(message, history):
+    # Format the prompt according to Gemma 2's template
+    prompt = f"<start_of_turn>user\n{message}<end_of_turn>\n<start_of_turn>model"
+    
+    # Tokenize input
+    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+    
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    # Decode the response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the model's response
+    if "<start_of_turn>model" in response:
+        return response.split("<start_of_turn>model")[-1].strip()
+    return response
+
+# Create the chat interface
+demo = gr.ChatInterface(
+    fn=chat_response,
+    title="Hausa AI Assistant",
+    description="A simple AI assistant powered by Gemma 2 2B"
+)
+
+# Launch the application
+if __name__ == "__main__":
+    demo.launch(server_name="0.0.0.0", server_port=7860)