modules/llm.py

from transformers import pipeline
import spaces

# Load a smaller Hugging Face instruct model
# Note: device is set to 'cuda' for ZeroGPU compatibility
llm_pipeline = pipeline(
    "text-generation", 
    model="mistralai/Mistral-7B-Instruct-v0.2", 
    device="cuda",
    # Using max_new_tokens instead of max_length to avoid the length validation error
    max_new_tokens=512,
    # Add truncation to handle long inputs
    truncation=True
)

@spaces.GPU(duration_s=60)  # Specify GPU duration for ZeroGPU
def ask_llm(prompt):
    # Create a more constrained prompt that instructs the model to only use provided context
    constrained_prompt = f"""<s>[INST] You are a helpful assistant that only answers based on the information provided in the context. 
If the information needed to answer the question is not in the context, say "I don't have enough information to answer this question."
Do not use any external knowledge or make assumptions beyond what is explicitly stated in the context.

{prompt} [/INST]</s>"""

    response = llm_pipeline(
        constrained_prompt, 
        return_full_text=False,
        do_sample=True,
        temperature=0.1,  # Lower temperature for more deterministic responses
        top_p=0.85,  # Slightly lower top_p to reduce creative variations
        repetition_penalty=1.2  # Discourage repetition
    )
    return response[0]["generated_text"]