Update app.py

app.py

@@ -1,37 +1,23 @@
 import os
 import gradio as gr
-from transformers import pipeline
-# from huggingface_hub import login
+import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# # Get the Hugging Face token from environment variables
-# HF_TOKEN = os.getenv('HF')
-
-# if not HF_TOKEN:
-#     raise ValueError("The HF environment variable is not set. Please set it to your Hugging Face token.")
-
-# # Authenticate with Hugging Face and save the token to the Git credentials helper
-# login(HF_TOKEN, add_to_git_credential=True)
-
-# Create the pipeline for text generation using the specified model
-# pipe = pipeline("text-generation", model="distilbert/distilgpt2", token=HF_TOKEN)
-pipe = pipeline("text-generation", model="openai-community/gpt2-medium")
+# Load the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2-medium")
+model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2-medium")
 
 # Define the initial prompt for the system
 system_prompt = """
 You are an AI model designed to provide concise information about big data analytics across various fields without mentioning the question. Respond with a focused, one-line answer that captures the essence of the key risk, benefit, or trend associated with the topic.
-
 input: What do you consider the most significant risk of over-reliance on big data analytics in stock market risk management?
 output: Increased market volatility.
-
 input: What is a major benefit of big data analytics in healthcare?
 output: Enhanced patient care through personalized treatment.
-
 input: What is a key challenge of big data analytics in retail?
 output: Maintaining data privacy and security.
-
 input: What is a primary advantage of big data analytics in manufacturing?
 output: Improved production efficiency and predictive maintenance.
-
 input: What is a significant risk associated with big data analytics in education?
 output: Potential widening of the achievement gap if data is not used equitably.
 """
@@ -41,9 +27,14 @@ def generate(text):
         # Combine the system prompt with the user's input
         prompt = system_prompt + f"\ninput: {text}\noutput:"
 
-        # Generate the response using the pipeline
-        responses = pipe(prompt, max_length=1024, num_return_sequences=1)
-        response_text = responses[0]['generated_text'].split("output:")[-1].strip()
+        # Tokenize the input
+        inputs = tokenizer(prompt, return_tensors="pt")
+
+        # Generate the response
+        outputs = model.generate(inputs["input_ids"], max_length=256)
+
+        # Convert the output to text
+        response_text = tokenizer.decode(outputs[0], skip_special_tokens=True).split("output:")[-1].strip()
 
         return response_text if response_text else "No valid response generated."
     
@@ -63,4 +54,4 @@ def launch_custom_interface():
     iface.launch()
 
 if __name__ == "__main__":
-    launch_custom_interface()
+    launch_custom_interface()