Spaces:
Running
Running
File size: 3,635 Bytes
cb3e313 c247815 eebf9a7 64fde65 c247815 cb3e313 297702e 8b78b70 cb3e313 a4a81f7 8b78b70 f24da04 eebf9a7 64fde65 ef17b91 64fde65 deed229 64fde65 deed229 64fde65 deed229 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
from transformers import pipeline
from huggingface_hub import login
import spaces
import json
import os
__export__ = ["GemmaLLM"]
class GemmaLLM:
def __init__(self):
login(token=os.environ.get("GEMMA_TOKEN"))
# 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:
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[2]["content"]
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:
message = self._get_questions_message(summary, num_questions, difficulty)
questions = self.generate(message)
return json.loads(questions[2]["content"].strip("```").replace("json\n", ""))
|