Intrigrate Google API

app.py

@@ -11,48 +11,9 @@ Implements the 5-module architecture from the design document:
  
 import os
 import re
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
- 
-# ---------------------------------------------------------------------------
-# MODULE 4 — Generation Module (LLM setup)
-# FIX: Correct model ID + read HF_TOKEN from environment (set as Space secret)
-# ---------------------------------------------------------------------------
-MODEL_ID  = "google/gemma-3-4b-it"   # smallest Gemma-3 — works on free tier
-HF_TOKEN  = os.environ.get("HF_TOKEN")  # set this in Space Settings → Secrets
- 
-_tokenizer = None
-_pipe      = None
- 
-def get_tokenizer():
-    """Returns the tokenizer, loading it on first call."""
-    global _tokenizer
-    if _tokenizer is None:
-        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN)
-    return _tokenizer
- 
-def get_pipeline():
-    """Returns the generation pipeline, loading model on first call (lazy load)."""
-    global _pipe
-    if _pipe is None:
-        tokenizer = get_tokenizer()
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_ID,
-            token=HF_TOKEN,
-            torch_dtype=torch.bfloat16,
-            device_map="auto",
-        )
-        _pipe = pipeline(
-            "text-generation",
-            model=model,
-            tokenizer=tokenizer,
-            max_new_tokens=512,
-            do_sample=False,
-        )
-    return _pipe
- 
- 
+from google import genai
+
 # ---------------------------------------------------------------------------
 # MODULE 3 — Pre-processing Module (6-Element Framework)
 # ---------------------------------------------------------------------------
@@ -166,10 +127,26 @@ def format_for_display(parsed: dict) -> tuple[str, str, str]:
 # ---------------------------------------------------------------------------
 # MODULE 4 — Generation Module (inference call)
 # ---------------------------------------------------------------------------
-def generate(prompt: str) -> str:
-    outputs = get_pipeline()(prompt, return_full_text=False)
-    return outputs[0]["generated_text"]
- 
+MODEL_ID= "gemma-4-31b-it" # LLM Model
+GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY") # 1. Fetch the API Key from the environment (configured in Settings -> Secrets)
+
+def generate_response(prompt):
+    # Check if the API Key is provided
+    if not GOOGLE_API_KEY:
+        return "⚠️ Error: API Key not found! Please configure GOOGLE_API_KEY in the Settings -> Variables and secrets page of your Space."
+    
+    try:
+        # 2. Initialise the Client using the new SDK
+        client = genai.Client(api_key=GOOGLE_API_KEY)
+        
+        # 3. Send the prompt to Google AI Studio for processing
+        response = client.models.generate_content(
+            model=MODEL_ID,
+            contents=prompt
+        )
+        return response.text
+    except Exception as e:
+        return f"❌ An error occurred: {str(e)}"
  
 # ---------------------------------------------------------------------------
 # MODULE 2 — Validation & Flow Management Module
@@ -202,7 +179,7 @@ def validate_and_evaluate(description: str, code: str):
  
     # --- Generation ---
     try:
-        raw_output = generate(prompt)
+        raw_output = generate_response(prompt)
     except Exception as exc:
         return "", "", "", f"❌ Model error: {exc}"
  

requirements.txt

@@ -3,4 +3,5 @@ torch>=2.2.0
 accelerate>=0.30.0
 gradio>=4.31.0
 sentencepiece
-protobuf
+protobuf
+google-genai