Update app.py

app.py

@@ -6,6 +6,40 @@ import requests
 import os 
 
 HF_API = os.getenv('HF_API')
+
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen-7B-Chat")
+model = AutoModelForCausalLM.from_pretrained("Qwen/Qwen-7B-Chat")
+
+def generate_summary(prompt):
+    # Add instructions to the prompt to signal that you want a summary
+    instructions = "Summarize the following text:"
+    prompt_with_instructions = f"{instructions}\n{prompt}"
+
+    # Tokenize the prompt text and return PyTorch tensors
+    inputs = tokenizer.encode(prompt_with_instructions, return_tensors="pt")
+
+    # Generate a response using the model
+    outputs = model.generate(inputs, max_length=512, num_return_sequences=1, pad_token_id=tokenizer.eos_token_id)
+
+    # Decode the response
+    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return summary
+
+def generate_response(prompt):
+    # Tokenize the prompt text and return PyTorch tensors
+    inputs = tokenizer.encode(prompt, return_tensors="pt")
+
+    # Generate a response using the model
+    outputs = model.generate(inputs, max_length=512, num_return_sequences=1)
+
+    # Decode the response
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response
+    
 # Function to search PubMed for articles
 def search_pubmed(query, retmax):
     Entrez.email = '[email protected]'
@@ -29,7 +63,7 @@ def search_pubmed(query, retmax):
     return pd.DataFrame(article_list)
 
 # Function to summarize articles using Hugging Face's API
-def summarize_with_huggingface(model, selected_articles):
+def summarize_with_huggingface(model, selected_articles, USE_LOCAL=True):
     API_URL = f"https://api-inference.huggingface.co/models/{model}"
     # Your Hugging Face API key
     API_KEY = HF_API 
@@ -46,9 +80,13 @@ def summarize_with_huggingface(model, selected_articles):
         "inputs": text_to_summarize,
         "parameters": {"max_length": 300}  # Adjust as needed
     }
-    # Make the POST request to the Hugging Face API
-    response = requests.post(API_URL, headers=headers, json=payload)
-    response.raise_for_status()  # Raise an HTTPError if the HTTP request returned an unsuccessful status code
+
+    if USE_LOCAL:
+        response = generate_response(text_to_summarize)
+    else:
+        # Make the POST request to the Hugging Face API
+        response = requests.post(API_URL, headers=headers, json=payload)
+        response.raise_for_status()  # Raise an HTTPError if the HTTP request returned an unsuccessful status code
     # The API returns a list of dictionaries. We extract the summary from the first one.
     return response.json()[0]['generated_text']