gemma.py

from dataset import prompts
import threading
import csv
import os
import time
from ollama import chat
from ollama import ChatResponse

NUM_THREADS = 10

SYSTEM_PROMPT = "Enhance the prompt to minimize token but should get fine results, each inference should generate 2 separate response, the question can be of coding as well so the out put should be Code First and explanation next the explanation must be neat and clear, if python code make sure the indentation should exist."

STOP_FLAG = False

def call_ollama(prompt):
    global STOP_FLAG
    if STOP_FLAG:
        return None
    try:
        response: ChatResponse = chat(model='gemma3:4b', messages=[
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": prompt}
        ])
        if hasattr(response, 'message'):
            return response.message.content if hasattr(response.message, 'content') else response.message['content']
        else:
            return response['message']['content']
    except Exception as e:
        print(f"Error: {e}")
        return "ERROR"


def get_log_path(thread_id):
    return f"thread_{thread_id}.log"


def read_progress(thread_id):
    path = get_log_path(thread_id)
    if not os.path.exists(path):
        return 0
    with open(path, "r") as f:
        try:
            return int(f.read().strip())
        except:
            return 0


def write_progress(thread_id, idx):
    with open(get_log_path(thread_id), "w") as f:
        f.write(str(idx))


def worker(thread_id, data, start_idx, end_idx):
    global STOP_FLAG
    filename = f"thread_{thread_id}_Question_{start_idx}_{end_idx}.csv"
    progress = read_progress(thread_id)

    mode = "a" if os.path.exists(filename) else "w"

    with open(filename, mode, newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if mode == "w":
            writer.writerow(["Question", "Answer 1", "Answer 2"])

        for i, q in enumerate(data):
            if i < progress:
                continue
            if STOP_FLAG:
                break
            res1 = call_ollama(q)
            if STOP_FLAG:
                break
            res2 = call_ollama(q)
            if STOP_FLAG:
                break
            writer.writerow([q, res1, res2])
            write_progress(thread_id, i + 1)
            time.sleep(0.2)


threads = []
chunk_size = len(prompts) // NUM_THREADS

for i in range(NUM_THREADS):
    start = i * chunk_size
    end = (i + 1) * chunk_size if i != NUM_THREADS - 1 else len(prompts)
    t = threading.Thread(target=worker, args=(i + 1, prompts[start:end], start, end))
    threads.append(t)
    t.start()

for t in threads:
    t.join()