Gemma.py

from transformers import pipeline
from huggingface_hub import login
import spaces
import json
import os
import openai

__export__ = ["GemmaLLM"]

class GemmaLLM:
    def __init__(self):
        login(token=os.environ.get("GEMMA_TOKEN"))
        openai.api_key = os.environ.get("OPENAI_API_KEY")
        
        # quant_config = quantization_config.BitsAndBytesConfig(
        #     load_in_8bit=True,
        #     llm_int8_threshold=6.0,
        #     llm_int8_has_fp16_weight=False,
        # )
        
        # model_id = "google/gemma-3-4b-it"
        # model_id = "google/gemma-3n-E4B-it-litert-preview"
        # model = AutoModelForCausalLM.from_pretrained(model_id, quantization_config=quant_config)
        # tokenizer = AutoTokenizer.from_pretrained(model_id)
        
        # self.model = pipeline("text-generation", model=model_id)
        
    @spaces.GPU
    def generate(self, message) -> str:
        response = openai.chat.completions.create(
            model="gpt-4.1",  # You can use "gpt-4" if you have access
            messages=message)
        return response.choices[0].message.content
        # outputs = self.model(message, max_new_tokens=1024)[0]["generated_text"]
        # return outputs

    def _get_summary_message(self, article, num_paragraphs) -> dict:
        summarize = "You are a helpful assistant. Your main task is to summarize articles. You will be given an article that you will generate a summary for. The summary should include all the key points of the article. ONLY RESPOND WITH THE SUMMARY!!!"

        summary = f"Summarize the data in the following JSON into {num_paragraphs} paragraph(s) so that it is easy to read and understand:\n"

        message = [{"role": "system", "content": [{"type": "text", "text": summarize}]},
                   {"role": "user",   "content": [{"type": "text", "text": summary + json.dumps(article, indent=4)}]}]

        return message

    def get_summary(self, article, num_paragraphs) -> str:
        message = self._get_summary_message(article, num_paragraphs)
        summary = self.generate(message)
        return summary

    def _get_questions_message(self, summary, num_questions, difficulty) -> dict:
        question = f"""
            You are a helpful assistant. Your main task is to generate {num_questions} multiple choice questions from an article. Respond in the following JSON structure and schema:\n\njson\n```{json.dumps(list((
            dict(question=str.__name__, correct_answer=str.__name__, false_answers=[str.__name__, str.__name__, str.__name__]),
            dict(question=str.__name__, correct_answer=str.__name__, false_answers=[str.__name__, str.__name__, str.__name__]),
            dict(question=str.__name__, correct_answer=str.__name__, false_answers=[str.__name__, str.__name__, str.__name__]))), indent=4)}```\n\nThere should only be {num_questions} questions generated. Each question should only have 3 false answers and 1 correct answer. The correct answer should be the most relevant answer based on the context derived from the article. False answers should not contain the correct answer. False answers should contain false information but also be reasonably plausible for answering the question. ONLY RESPOND WITH RAW JSON!!!
        """

        questions = f"Generate {difficulty.lower()} questions and answers from the following article:\n"

        message = [{"role": "system", "content": [{"type": "text", "text": question}]},
                   {"role": "user", "content": [{"type": "text", "text": questions + summary}]}]

        return message
    
    def get_questions(self, summary, num_questions, difficulty) -> dict:
        # print("Getting questions message...")
        message = self._get_questions_message(summary, num_questions, difficulty)
        # print("Generating questions...")
        questions = self.generate(message)
        # print(questions)
        questions = questions.strip("```")
        questions = questions.replace("json\n", "")
        questions = json.loads(questions)
        print(questions)
        return questions

if __name__ == "__main__":
    gemma = GemmaLLM()
    summary = gemma.get_summary('''
    Iran could reach an agreement similar to the 2015 nuclear deal to end its current conflict with Israel, the Iranian foreign minister said Saturday, according to state media.

    “It is clear the Israelis are against diplomacy and do not want the current crisis to be resolved peacefully,” Abbas Araghchi said, according to state-run news agency IRNA.

    “We are, however, certainly prepared to reach a negotiated solution, similar to the one we reached in 2015,” the foreign minister said, according to IRNA.

    Remember: The 2015 deal, formally known as the Joint Comprehensive Plan of Action, required Iran to limit its nuclear program in return for provisions including relief on sanctions. US President Donald Trump withdrew from the landmark agreement in 2018 and Iran has increasingly grown its nuclear program since then.

    More from the foreign minister: “That agreement was the result of two years of tireless negotiations, and when it was signed, it was welcomed by the entire world,” Araghchi said of the 2015 deal, according to IRNA.

    “Diplomacy can be effective — just as it was effective in the past — and can be effective in the future as well. But to return to the path of diplomacy, the aggression must be halted,” he continued, according to the news agency.
    ''', 1)
    print(summary)
    questions = gemma.get_questions(summary, 3, "Easy")
    print(json.dumps(questions, indent=4))