Update space

Gemma_Model.py

@@ -0,0 +1,102 @@
+from transformers import AutoTokenizer, BitsAndBytesConfig, Gemma3ForCausalLM
+import torch
+import json
+
+__export__ = ["GemmaLLM"]
+
+class GemmaLLM:
+    
+    def __init__(self):
+        model_id = "google/gemma-3-1b-it"
+        quantization_config = BitsAndBytesConfig(load_in_8bit=True)
+        
+        self.model = Gemma3ForCausalLM.from_pretrained(
+            model_id,
+            device_map="cpu",
+            quantization_config=quantization_config,
+            low_cpu_mem_usage=True,
+            torch_dtype=torch.float16,
+        ).eval()
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+    def generate(self, message) -> str:
+        print("Generating...")
+        inputs = self.tokenizer.apply_chat_template(
+            message,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+        ).to(self.model.device)
+
+        input_length = inputs["input_ids"].shape[1]
+
+        with torch.inference_mode():
+            outputs = self.model.generate(
+                **inputs, max_new_tokens=1024,
+            )[0][input_length:]
+
+        outputs = self.tokenizer.decode(outputs, skip_special_tokens=True)
+        
+        print("Completed generating!")
+        
+        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, message) -> str:
+        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} 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, message) -> dict:
+        questions = self.generate(message)
+
+        return json.loads(questions.strip("```").replace("json\n", ""))

News.py

@@ -0,0 +1,62 @@
+from newsapi import NewsApiClient
+from newspaper import Article
+import os
+
+__export__ = ["News"]
+
+class News:
+    __EX_SOURCES__ = {"ABC News", "Bloomberg", "The Hill", "Fox Sports", "Google News"}
+    __CATEGORIES__ = {
+        "business",
+        "entertainment",
+        "general",
+        "health",
+        "science",
+        "sports",
+        "technology"
+    }
+    
+    def __init__(self):
+        newsapi_key = os.environ.get("NEWS_API_KEY")
+        self.newsapi = NewsApiClient(api_key=newsapi_key)
+
+    def get_sources(self, category=None):
+        sources = self.newsapi.get_sources(language="en", country="us", category=category)["sources"]
+        sources = {source["name"] for source in sources if source["name"] not in self.__EX_SOURCES__}
+        print(sources)
+        return sources
+        
+
+    def get_top_headlines(self, num_headlines=None, category=None):
+        sources = self.get_sources(category=category)
+        
+        headlines = self.newsapi.get_top_headlines(
+            sources=", ".join(sources),
+            page_size=num_headlines
+        )["articles"]
+        
+        return headlines
+
+    def get_headlines(self, num_headlines=None, query=None):
+        sources = self.get_sources()
+        
+        headlines = self.newsapi.get_everything(
+            q=query,
+            sources=", ".join(sources),
+            page_size=num_headlines
+        )["articles"]
+        
+        return headlines
+
+    def get_articles_from_headlines(self, headlines):
+        for headline in headlines:
+            article = Article(headline["url"])
+            article.download()
+            article.parse()
+            headline["content"] = article.text
+            # headline["authors"] = article.authors
+            headline["source"] = headline["source"]["name"]
+            del headline["author"]
+            # headline.pop("author", None)
+
+        return headlines