edited generate code

app.py

@@ -16,15 +16,6 @@ pipe = pipeline("text2text-generation", model="google/flan-t5-small")
 categories = ('Heart', 'Oblong', 'Oval', 'Round', 'Square')
 learn = load_learner('model.pkl')
 
-# Initialize the Code Llama Instruct pipeline (example with 7B model)
-llama_model_id = "meta-llama/CodeLlama-7b-Instruct-hf"
-llama_pipeline = pipeline(
-    "text-generation",
-    model=llama_model_id,
-    model_kwargs={"torch_dtype": torch.bfloat16},
-    device_map="auto"
-)
-
 # Überprüfe, ob das Zugriffstoken vorhanden ist
 if access_token is None:
     raise ValueError("Access token is missing. Make sure it is set as an environment variable.")
@@ -69,36 +60,41 @@ async def face_analyse(file: UploadFile = File(...)):
     # Assuming categories is a list of category labels
     return dict(zip(categories, map(float, probs)))
 
-@app.post("/extract-frame-details")
-def extract_frame_details(text: str):
+# Initialisiere das Modell und den Tokenizer
+model = "meta-llama/CodeLlama-7b-hf"
+tokenizer = AutoTokenizer.from_pretrained(model)
+llama_pipeline = pipeline(
+    "text-generation",
+    model=model,
+    torch_dtype=torch.float16,
+    device_map="auto",
+)
+
+@app.get("/generate_json")
+def generate_code(text: str):
     """
-    Using the Code Llama Instruct pipeline from `transformers`, extract frame
-    details from the given input text. The model used is `meta-llama/CodeLlama-7b-Instruct-hf`.
+    Using the Code Llama pipeline from `transformers`, generate code
+    from the given input text. The model used is `meta-llama/CodeLlama-7b-hf`.
     """
-    messages = [
-        {"role": "system", "content": "Please provide details about frames in JSON format."},
-        {"role": "user", "content": text},
-    ]
+    try:
+        sequences = llama_pipeline(
+            text,
+            do_sample=True,
+            top_k=10,
+            temperature=0.1,
+            top_p=0.95,
+            num_return_sequences=1,
+            eos_token_id=tokenizer.eos_token_id,
+            max_length=200,
+        )
 
-    terminators = [
-        llama_pipeline.tokenizer.eos_token_id,
-        llama_pipeline.tokenizer.convert_tokens_to_ids("")
-    ]
+        generated_text = sequences[0]["generated_text"]
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-    outputs = llama_pipeline(
-        messages,
-        max_new_tokens=256,
-        eos_token_id=terminators,
-        do_sample=True,
-        temperature=0.6,
-        top_p=0.9,
-    )
-    
-    generated_text = outputs[0]["generated_text"]
-    
-    try:
-        extracted_info = json.loads(generated_text)
-    except json.JSONDecodeError:
-        return {"error": "Failed to parse the generated text as JSON."}
+    return {"generated_text": generated_text}
 
-    return extracted_info
+# Beispielaufruf mit curl:
+# curl -X 'GET' \
+#   'http://localhost:8000/generate_code?text=import%20socket%0A%0Adef%20ping_exponential_backoff(host%3A%20str)%3A' \
+#   -H 'accept: application/json'