Update generate.py

generate.py

@@ -1,38 +1,59 @@
 import os
-
-import openai
 from dotenv import load_dotenv
 
-# Load environment variables from .env file
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+
+# Load environment variables if needed
 load_dotenv()
 
-# Retrieve the OpenAI API key from the environment variable
-openai.api_key = os.getenv('api_key')
+# Use the Qwen2.5-7B-Instruct-1M model from Hugging Face
+MODEL_NAME = "Qwen/Qwen2.5-7B-Instruct-1M"
 
-def generate_response(retrieved_texts, query, max_tokens=2000):
-    """
-    Generates a response based on the retrieved texts and query.
+# Initialize tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",        # or "cpu", "cuda", etc. as appropriate
+    trust_remote_code=True
+)
 
+# Create pipeline
+qwen_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer
+)
+
+def generate_response(retrieved_texts, query, max_new_tokens=512):
+    """
+    Generates a response based on the retrieved texts and query using the Qwen pipeline.
     Args:
-    retrieved_texts (list): List of retrieved text strings.
-    query (str): Query string.
-    max_tokens (int): Maximum number of tokens for the response.
+        retrieved_texts (list): List of retrieved text strings.
+        query (str): The user's query string.
+        max_new_tokens (int): Maximum number of tokens for the generated answer.
 
     Returns:
-    str: Generated response.
+        str: Generated response.
     """
+    # Construct a simple prompt using your retrieved context
     context = "\n".join(retrieved_texts)
-    prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
-    
-    response = openai.ChatCompletion.create(
-        model="gpt-3.5-turbo",
-        messages=[
-            {"role": "system", "content": "You are a helpful assistant."},
-            {"role": "user", "content": prompt}
-        ],
-        max_tokens=max_tokens,
-        n=1,
-        stop=None,
-        temperature=0.5,
+    prompt = f"Context:\n{context}\n\nQuestion: {query}\nAnswer:"
+
+    # Generate the text
+    result = qwen_pipeline(
+        prompt,
+        max_new_tokens=max_new_tokens,
+        do_sample=True,      # or False if you prefer deterministic output
+        temperature=0.7,     # adjust as needed
     )
-    return response.choices[0].message['content']
+
+    # Extract the generated text from the pipeline's output
+    generated_text = result[0]["generated_text"]
+
+    # Optional: Clean up the output to isolate the answer portion
+    if "Answer:" in generated_text:
+        answer_part = generated_text.split("Answer:")[-1].strip()
+    else:
+        answer_part = generated_text
+
+    return answer_part