[Gemma] Google Gemma3와 Flask 서버 연동해서 사용하기

모델 연동 및 Flask 설치

앞서 적었던 포스트대로 먼저 Gemma 모델을 불러와준다.
그 후, 가상환경을 실행시켜 Flask 서버를 설치해준다.

pip install Flask

코드 작성

from flask import Flask, request, Response
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
import json
import re

app = Flask(__name__)

model_id = "google/gemma-3-1b-it"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto", torch_dtype="auto")
pipe = pipeline("text-generation", model=model, tokenizer=tokenizer)

GARBAGE_GUIDE = {
   # 필요한 항목 추가
}

## 지정한 경로로 요청 조냄
@app.route("/generate", methods=["POST"])
def generate():
    data = request.json
    materials = data.get("material", [])
    if isinstance(materials, str):
        materials = [materials]

    explanations = {}

    for material in materials:
        material = material.strip().upper()

        if material not in GARBAGE_GUIDE:
            explanations[material] = "❌ 지원하지 않는 분류입니다."
            continue

        guide_text = GARBAGE_GUIDE[material]
        prompt = f"""{guide_text}

이 내용을 사용자에게 친절하고 간결하게 설명해주세요.
다른 정보는 말하지 말고, 위 내용만 정리해서 말해주세요.
"""

        result = pipe(prompt, max_new_tokens=300, do_sample=False, return_full_text=False)[0]["generated_text"]

        # 반복되는 ** 기호 제거
        cleaned = re.sub(r"\*{2,}", "", result)
        answer_only = cleaned.strip()

        if not answer_only:
            answer_only = "⚠️ 설명 생성에 실패했습니다."

        explanations[material] = answer_only

    return Response(json.dumps({"responses": explanations}, ensure_ascii=False), mimetype="application/json")

if __name__ == "__main__":
    app.run(debug=True)

나는 /generate 경로로 보내줬지만 원하는 경로로 정보를 Flask 서버에 넘겨준다.
나는 욜로 모델에게 material 정보를 받아 Gemma로 넘겨주는 구조였기에, 테스트 구조 상에서는 material을 내가 임의로 결정해 넘겨주었다.

그 후, 파일을 실행시켜준다.

지정한 Flask 경로로 요청 보내기

• 터미널에서 바로 보내기

• 파이썬 코드를 통해 보내기

import requests

url = "http://127.0.0.1:5000/generate"
payload = {"material": ["PLASTIC", "CAN"]}

res = requests.post(url, json=payload)
print(res.json())

이렇게 요청을 보내고나면 터미널창에서 Gemma가 보내준 정보를 확인할 수 있다.
아직 다듬는 중이라서 정보가 뒤죽박죽으로 출력되었지만 ... 프롬프트 엔지니어링 방법에 대해 더 공부해봐야겠다.